fix(python/sdk): Fix mypy and pytest-httpx errors (#6609)

GitOrigin-RevId: 141b374033010e6154edd910e7082dba37159613

Author: ploeber

assemblyai/api.py

@@ -41,7 +41,7 @@ def create_transcript(
             by_alias=True,
         ),
     )
-    if response.status_code != httpx.codes.ok:
+    if response.status_code != httpx.codes.OK:
         raise types.TranscriptError(
             f"failed to transcribe url {request.audio_url}: {_get_error_message(response)}"
         )
@@ -57,7 +57,7 @@ def get_transcript(
         f"{ENDPOINT_TRANSCRIPT}/{transcript_id}",
     )
 
-    if response.status_code != httpx.codes.ok:
+    if response.status_code != httpx.codes.OK:
         raise types.TranscriptError(
             f"failed to retrieve transcript {transcript_id}: {_get_error_message(response)}",
         )
@@ -73,7 +73,7 @@ def delete_transcript(
         f"{ENDPOINT_TRANSCRIPT}/{transcript_id}",
     )
 
-    if response.status_code != httpx.codes.ok:
+    if response.status_code != httpx.codes.OK:
         raise types.TranscriptError(
             f"failed to delete transcript {transcript_id}: {_get_error_message(response)}",
         )
@@ -100,7 +100,7 @@ def upload_file(
         content=audio_file,
     )
 
-    if response.status_code != httpx.codes.ok:
+    if response.status_code != httpx.codes.OK:
         raise types.TranscriptError(
             f"Failed to upload audio file: {_get_error_message(response)}"
         )
@@ -125,7 +125,7 @@ def export_subtitles_srt(
         params=params,
     )
 
-    if response.status_code != httpx.codes.ok:
+    if response.status_code != httpx.codes.OK:
         raise types.TranscriptError(
             f"failed to export SRT for transcript {transcript_id}: {_get_error_message(response)}"
         )
@@ -150,7 +150,7 @@ def export_subtitles_vtt(
         params=params,
     )
 
-    if response.status_code != httpx.codes.ok:
+    if response.status_code != httpx.codes.OK:
         raise types.TranscriptError(
             f"failed to export VTT for transcript {transcript_id}: {_get_error_message(response)}"
         )
@@ -172,7 +172,7 @@ def word_search(
         ),
     )
 
-    if response.status_code != httpx.codes.ok:
+    if response.status_code != httpx.codes.OK:
         raise types.TranscriptError(
             f"failed to search words in transcript {transcript_id}: {_get_error_message(response)}"
         )
@@ -223,7 +223,7 @@ def get_sentences(
         f"{ENDPOINT_TRANSCRIPT}/{transcript_id}/sentences",
     )
 
-    if response.status_code != httpx.codes.ok:
+    if response.status_code != httpx.codes.OK:
         raise types.TranscriptError(
             f"failed to retrieve sentences for transcript {transcript_id}: {_get_error_message(response)}"
         )
@@ -239,7 +239,7 @@ def get_paragraphs(
         f"{ENDPOINT_TRANSCRIPT}/{transcript_id}/paragraphs",
     )
 
-    if response.status_code != httpx.codes.ok:
+    if response.status_code != httpx.codes.OK:
         raise types.TranscriptError(
             f"failed to retrieve paragraphs for transcript {transcript_id}: {_get_error_message(response)}"
         )
@@ -262,7 +262,7 @@ def list_transcripts(
         ),
     )
 
-    if response.status_code != httpx.codes.ok:
+    if response.status_code != httpx.codes.OK:
         raise types.AssemblyAIError(
             f"failed to retrieve transcripts: {_get_error_message(response)}"
         )
@@ -283,7 +283,7 @@ def lemur_question(
         timeout=http_timeout,
     )
 
-    if response.status_code != httpx.codes.ok:
+    if response.status_code != httpx.codes.OK:
         raise types.LemurError(
             f"failed to call Lemur questions: {_get_error_message(response)}"
         )
@@ -304,7 +304,7 @@ def lemur_summarize(
         timeout=http_timeout,
     )
 
-    if response.status_code != httpx.codes.ok:
+    if response.status_code != httpx.codes.OK:
         raise types.LemurError(
             f"failed to call Lemur summary: {_get_error_message(response)}"
         )
@@ -325,7 +325,7 @@ def lemur_action_items(
         timeout=http_timeout,
     )
 
-    if response.status_code != httpx.codes.ok:
+    if response.status_code != httpx.codes.OK:
         raise types.LemurError(
             f"failed to call Lemur action items: {_get_error_message(response)}"
         )
@@ -346,7 +346,7 @@ def lemur_task(
         timeout=http_timeout,
     )
 
-    if response.status_code != httpx.codes.ok:
+    if response.status_code != httpx.codes.OK:
         raise types.LemurError(
             f"failed to call Lemur task: {_get_error_message(response)}"
         )
@@ -358,13 +358,13 @@ def lemur_purge_request_data(
     client: httpx.Client,
     request: types.LemurPurgeRequest,
     http_timeout: Optional[float],
-) -> types.LemurPurgeRequest:
+) -> types.LemurPurgeResponse:
     response = client.delete(
         f"{ENDPOINT_LEMUR_BASE}/{request.request_id}",
         timeout=http_timeout,
     )
 
-    if response.status_code != httpx.codes.ok:
+    if response.status_code != httpx.codes.OK:
         raise types.LemurError(
             f"Failed to purge LeMUR request data for provided request ID: {request.request_id}. Error: {_get_error_message(response)}"
         )
@@ -374,7 +374,7 @@ def lemur_purge_request_data(
 
 def lemur_get_response_data(
     client: httpx.Client,
-    request_id: int,
+    request_id: str,
     http_timeout: Optional[float],
 ) -> Union[
     types.LemurStringResponse,
@@ -385,7 +385,7 @@ def lemur_get_response_data(
         timeout=http_timeout,
     )
 
-    if response.status_code != httpx.codes.ok:
+    if response.status_code != httpx.codes.OK:
         raise types.LemurError(
             f"Failed to get LeMUR response data for provided request ID: {request_id}. Error: {_get_error_message(response)}"
         )
@@ -409,7 +409,7 @@ def create_temporary_token(
         timeout=http_timeout,
     )
 
-    if response.status_code != httpx.codes.ok:
+    if response.status_code != httpx.codes.OK:
         raise types.AssemblyAIError(
             f"Failed to create temporary token: {_get_error_message(response)}"
         )

assemblyai/lemur.py

@@ -12,7 +12,7 @@ def __init__(
         self,
         *,
         client: _client.Client,
-        sources: List[types.LemurSource],
+        sources: Optional[List[types.LemurSource]],
     ) -> None:
         self._client = client
 

assemblyai/types.py

@@ -521,7 +521,7 @@ class RawTranscriptionConfig(BaseModel):
     iab_categories: Optional[bool]
     "Enable Topic Detection."
 
-    custom_spelling: Optional[List[Dict[str, List[str]]]]
+    custom_spelling: Optional[List[Dict[str, Union[str, List[str]]]]]
     "Customize how words are spelled and formatted using to and from values"
 
     disfluencies: Optional[bool]
@@ -649,10 +649,11 @@ def __init__(
             speech_threshold: Reject audio files that contain less than this fraction of speech. Valid values are in the range [0,1] inclusive.
             raw_transcription_config: Create the config from a `RawTranscriptionConfig`
         """
-        self._raw_transcription_config = raw_transcription_config
-
-        if raw_transcription_config is None:
-            self._raw_transcription_config = RawTranscriptionConfig()
+        self._raw_transcription_config = (
+            raw_transcription_config
+            if raw_transcription_config is not None
+            else RawTranscriptionConfig()
+        )
 
         # explicit configurations have higher priority if `raw_transcription_config` has been passed as well
         self.language_code = language_code
@@ -914,17 +915,21 @@ def iab_categories(self, enable: Optional[bool]) -> None:
         self._raw_transcription_config.iab_categories = enable
 
     @property
-    def custom_spelling(self) -> Optional[Dict[str, List[str]]]:
+    def custom_spelling(self) -> Optional[Dict[str, Union[str, List[str]]]]:
         "Returns the current set custom spellings."
 
         if self._raw_transcription_config.custom_spelling is None:
             return None
 
         custom_spellings = {}
         for custom_spelling in self._raw_transcription_config.custom_spelling:
-            custom_spellings[custom_spelling["from"]] = custom_spelling["to"]
+            _from = custom_spelling["from"]
+            if isinstance(_from, str):
+                custom_spellings[_from] = custom_spelling["to"]
+            else:
+                raise ValueError("`from` argument must be a string!")
 
-        return custom_spellings
+        return custom_spellings if custom_spelling else None
 
     @property
     def disfluencies(self) -> Optional[bool]:
@@ -938,8 +943,6 @@ def disfluencies(self, enable: Optional[bool]) -> None:
 
         self._raw_transcription_config.disfluencies = enable
 
-        return self
-
     @property
     def sentiment_analysis(self) -> Optional[bool]:
         "Returns the status of the Sentiment Analysis feature."
@@ -953,7 +956,7 @@ def sentiment_analysis(self, enable: Optional[bool]) -> None:
         self._raw_transcription_config.sentiment_analysis = enable
 
     @property
-    def auto_chapters(self) -> bool:
+    def auto_chapters(self) -> Optional[bool]:
         "Returns the status of the Auto Chapters feature."
 
         return self._raw_transcription_config.auto_chapters
@@ -971,7 +974,7 @@ def auto_chapters(self, enable: Optional[bool]) -> None:
         self._raw_transcription_config.auto_chapters = enable
 
     @property
-    def entity_detection(self) -> bool:
+    def entity_detection(self) -> Optional[bool]:
         "Returns whether Entity Detection feature is enabled or not."
 
         return self._raw_transcription_config.entity_detection
@@ -1076,7 +1079,7 @@ def set_casing_and_formatting(
 
     def set_speaker_diarization(
         self,
-        enable: bool = True,
+        enable: Optional[bool] = True,
         speakers_expected: Optional[int] = None,
     ) -> Self:
         """
@@ -1261,7 +1264,7 @@ def set_custom_spelling(
 
     def set_summarize(
         self,
-        enable: bool = True,
+        enable: Optional[bool] = True,
         model: Optional[SummarizationModel] = None,
         type: Optional[SummarizationType] = None,
     ) -> Self:
@@ -1866,13 +1869,6 @@ def source(self) -> Sourcable:
         """
         return self._source
 
-    @property
-    def context(self) -> Optional[Union[str, Dict[str, Any]]]:
-        """
-        An optional context on the source (can be a string or an arbitrary dictionary)
-        """
-        return self._context
-
     @property
     def type(self) -> LemurSourceType:
         """

tests/unit/test_auto_chapters.py

@@ -34,9 +34,6 @@ def test_auto_chapters_fails_without_punctuation(httpx_mock: HTTPXMock):
     # Check that the error was raised before any requests were made
     assert len(httpx_mock.get_requests()) == 0
 
-    # Inform httpx_mock that it's okay we didn't make any requests
-    httpx_mock.reset(False)
-
 
 def test_auto_chapters_disabled_by_default(httpx_mock: HTTPXMock):
     """

tests/unit/test_content_safety.py

@@ -229,6 +229,3 @@ def test_content_safety_with_invalid_confidence_threshold(
 
     # Check that the error was raised before any requests were made
     assert len(httpx_mock.get_requests()) == 0
-
-    # Inform httpx_mock that it's okay we didn't make any requests
-    httpx_mock.reset(False)

tests/unit/test_redact_pii.py

@@ -115,9 +115,6 @@ def test_redact_pii_fails_without_policies(httpx_mock: HTTPXMock):
     # Check that the error was raised before any requests were made
     assert len(httpx_mock.get_requests()) == 0
 
-    # Inform httpx_mock that it's okay we didn't make any requests
-    httpx_mock.reset(False)
-
 
 def test_redact_pii_params_excluded_when_disabled(httpx_mock: HTTPXMock):
     """

tests/unit/test_summarization.py

@@ -37,9 +37,6 @@ def test_summarization_fails_without_required_field(
     # Check that the error was raised before any requests were made
     assert len(httpx_mock.get_requests()) == 0
 
-    # Inform httpx_mock that it's okay we didn't make any requests
-    httpx_mock.reset(False)
-
 
 def test_summarization_disabled_by_default(httpx_mock: HTTPXMock):
     """