diff --git a/assemblyai/__version__.py b/assemblyai/__version__.py index 0a0a43a5..38cf6dbe 100644 --- a/assemblyai/__version__.py +++ b/assemblyai/__version__.py @@ -1 +1 @@ -__version__ = "1.9.0" +__version__ = "1.9.1" diff --git a/assemblyai/prerecorded/v2/async_client.py b/assemblyai/prerecorded/v2/async_client.py index 975468f5..0cb5f4f6 100644 --- a/assemblyai/prerecorded/v2/async_client.py +++ b/assemblyai/prerecorded/v2/async_client.py @@ -7,7 +7,7 @@ import stat from types import TracebackType from typing import ( - AsyncIterator, + AsyncGenerator, Awaitable, BinaryIO, Callable, @@ -69,7 +69,7 @@ def _peek_length(stream: BinaryIO) -> Optional[int]: async def _aiter_stream( stream: BinaryIO, chunk_size: int = _UPLOAD_CHUNK_SIZE, -) -> AsyncIterator[bytes]: +) -> AsyncGenerator[bytes, None]: """Yields `stream` in chunks. Each read runs off the event loop.""" while True: @@ -83,7 +83,7 @@ async def _aiter_stream( async def _aiter_path( path: str, chunk_size: int = _UPLOAD_CHUNK_SIZE, -) -> AsyncIterator[bytes]: +) -> AsyncGenerator[bytes, None]: """Yields the file at `path` in chunks. Opens and reads it off the loop.""" audio_file = await _run_in_thread(_open_binary, path, "rb") @@ -96,7 +96,7 @@ async def _aiter_path( def _upload_request( data: AudioSource, -) -> Tuple[Union[bytes, AsyncIterator[bytes]], Dict[str, str]]: +) -> Tuple[Union[bytes, AsyncGenerator[bytes, None]], Dict[str, str]]: """ Turns audio into an upload body and the headers that describe it. @@ -112,7 +112,7 @@ def _upload_request( if isinstance(data, (str, os.PathLike)): path = os.fspath(data) - content: Union[bytes, AsyncIterator[bytes]] = _aiter_path(path) + content: Union[bytes, AsyncGenerator[bytes, None]] = _aiter_path(path) try: length: Optional[int] = os.path.getsize(path) except OSError: @@ -221,11 +221,15 @@ async def upload_file(self, data: AudioSource) -> str: content, headers = _upload_request(data) - return await async_api.upload_file( - self._client.http_client, - content, - headers=headers, - ) + try: + return await async_api.upload_file( + self._client.http_client, + content, + headers=headers, + ) + finally: + if not isinstance(content, bytes): + await content.aclose() async def submit( self, diff --git a/tests/unit/test_async_transcriber.py b/tests/unit/test_async_transcriber.py index bdcd683c..dde24b12 100644 --- a/tests/unit/test_async_transcriber.py +++ b/tests/unit/test_async_transcriber.py @@ -915,3 +915,67 @@ async def test_canonical_module_path_exports_both_transcribers(): assert AsyncTranscriber is aai.AsyncTranscriber assert Transcriber is aai.Transcriber + + +@pytest.mark.parametrize("outcome", ["write_error", "cancelled", "early_response"]) +async def test_interrupted_path_upload_closes_file_before_returning( + tmp_path, monkeypatch, outcome +): + from assemblyai.prerecorded.v2 import async_client as transcriber_module + + path = tmp_path / "audio.wav" + path.write_bytes(b"audio" * 300000) + opened = [] + real_open = transcriber_module._open_binary + + def track_open(*args): + stream = real_open(*args) + opened.append(stream) + return stream + + class FailingTransport(httpx.AsyncBaseTransport): + async def handle_async_request(self, request): + async for chunk in request.stream: + assert chunk + if outcome == "cancelled": + raise asyncio.CancelledError() + if outcome == "early_response": + return httpx.Response( + 200, json={"upload_url": "https://example.org/audio.wav"} + ) + raise httpx.WriteError("upload interrupted", request=request) + + monkeypatch.setattr(transcriber_module, "_open_binary", track_open) + async with aai.AsyncClient() as client: + await client.http_client.aclose() + client._http_client = httpx.AsyncClient( + base_url=aai.settings.base_url, transport=FailingTransport() + ) + transcriber = aai.AsyncTranscriber(client=client) + if outcome == "early_response": + assert ( + await transcriber.upload_file(path) == "https://example.org/audio.wav" + ) + else: + error = ( + asyncio.CancelledError if outcome == "cancelled" else httpx.WriteError + ) + with pytest.raises(error): + await transcriber.upload_file(path) + assert len(opened) == 1 + assert opened[0].closed + + +async def test_upload_error_keeps_caller_file_open(monkeypatch): + stream = io.BytesIO(b"audio") + + async def fail_upload(client, content, headers): + await content.__anext__() + raise httpx.WriteError("upload interrupted") + + monkeypatch.setattr(async_api, "upload_file", fail_upload) + async with aai.AsyncTranscriber() as transcriber: + with pytest.raises(httpx.WriteError): + await transcriber.upload_file(stream) + assert not stream.closed + stream.close()