fix: close interrupted asynchronous file uploads

assemblyai/__version__.py

@@ -1 +1 @@
-__version__ = "1.8.2"
+__version__ = "1.8.3"

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,

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()