fix: size async uploads from logical stream positions

assemblyai/prerecorded/v2/async_client.py

@@ -4,7 +4,6 @@
 
 import asyncio
 import os
-import stat
 from types import TracebackType
 from typing import (
     AsyncIterator,
@@ -39,29 +38,18 @@
 
 def _peek_length(stream: BinaryIO) -> Optional[int]:
     """
-    Returns the bytes left in `stream`, or `None` if the size is unknown.
+    Returns the logical bytes left in `stream`, or `None` if the size is unknown.
 
     A known size lets the upload send `Content-Length`. Pipes and sockets
     cannot report a size.
     """
 
     try:
         offset = stream.tell()
+        length = stream.seek(0, os.SEEK_END)
+        stream.seek(offset)
     except (AttributeError, OSError, ValueError):
-        offset = 0
-
-    try:
-        stat_result = os.fstat(stream.fileno())
-        if not stat.S_ISREG(stat_result.st_mode):
-            # A pipe or terminal reports st_size 0, which is not its length.
-            raise OSError
-        length = stat_result.st_size
-    except (AttributeError, OSError, ValueError):
-        try:
-            length = stream.seek(0, os.SEEK_END)
-            stream.seek(offset)
-        except (AttributeError, OSError, ValueError):
-            return None
+        return None
 
     return max(length - offset, 0)
 
@@ -219,7 +207,7 @@ async def upload_file(self, data: AudioSource) -> str:
         Returns: The URL of the uploaded audio file.
         """
 
-        content, headers = _upload_request(data)
+        content, headers = await _run_in_thread(_upload_request, data)
 
         return await async_api.upload_file(
             self._client.http_client,

tests/unit/test_async_transcriber.py

@@ -1,8 +1,12 @@
 import asyncio
+import bz2
+import gzip
 import io
 import json
+import lzma
 import os
 import re
+import threading
 from urllib.parse import urlencode
 
 import httpx
@@ -915,3 +919,59 @@ async def test_canonical_module_path_exports_both_transcribers():
 
     assert AsyncTranscriber is aai.AsyncTranscriber
     assert Transcriber is aai.Transcriber
+
+
+@pytest.mark.parametrize("open_compressed", [gzip.open, bz2.open, lzma.open])
+@pytest.mark.parametrize("offset", [0, 128])
+async def test_upload_file_uses_logical_stream_length(
+    httpx_mock: HTTPXMock, tmp_path, open_compressed, offset
+):
+    audio = b"audio" * 1024
+    path = tmp_path / "audio.compressed"
+    with open_compressed(path, "wb") as compressed:
+        compressed.write(audio)
+    httpx_mock.add_response(
+        url=UPLOAD_URL,
+        method="POST",
+        json={"upload_url": "https://example.org/uploaded.wav"},
+    )
+
+    with open_compressed(path, "rb") as stream:
+        stream.read(offset)
+        async with aai.AsyncTranscriber() as transcriber:
+            await transcriber.upload_file(stream)
+        assert not stream.closed
+
+    request = httpx_mock.get_requests()[0]
+    assert request.read() == audio[offset:]
+    assert int(request.headers["content-length"]) == len(request.read())
+
+
+async def test_upload_file_size_probe_does_not_block_event_loop(httpx_mock: HTTPXMock):
+    started = threading.Event()
+    resume = threading.Event()
+
+    class PausedSizeProbe(io.BytesIO):
+        def seek(self, offset, whence=os.SEEK_SET):
+            if whence == os.SEEK_END:
+                started.set()
+                assert resume.wait(timeout=2)
+            return super().seek(offset, whence)
+
+    httpx_mock.add_response(
+        url=UPLOAD_URL,
+        method="POST",
+        json={"upload_url": "https://example.org/uploaded.wav"},
+    )
+    async with aai.AsyncTranscriber() as transcriber:
+        uploading = asyncio.create_task(
+            transcriber.upload_file(PausedSizeProbe(b"audio"))
+        )
+        try:
+            loop = asyncio.get_running_loop()
+            assert await asyncio.wait_for(
+                loop.run_in_executor(None, started.wait, 1), timeout=2
+            )
+        finally:
+            resume.set()
+            await uploading