Cap upload stream reads to remaining file bytes

mapillary_tools/upload_api_v4.py

@@ -71,11 +71,35 @@ def chunkize_byte_stream(
             raise ValueError("Expect positive chunk size")
 
         while True:
-            data = stream.read(chunk_size)
+            to_read = cls._bounded_read_size(stream, chunk_size)
+            if to_read == 0:
+                break
+            data = stream.read(to_read)
             if not data:
                 break
             yield data
 
+    @classmethod
+    def _bounded_read_size(cls, stream: T.IO[bytes], chunk_size: int) -> int:
+        """Cap read() to remaining bytes so FileIO does not allocate chunk_size (e.g. 1 GiB)."""
+        remaining = cls._remaining_bytes(stream)
+        if remaining is None:
+            return chunk_size
+        if remaining <= 0:
+            return 0
+        return min(chunk_size, remaining)
+
+    @classmethod
+    def _remaining_bytes(cls, stream: T.IO[bytes]) -> int | None:
+        try:
+            pos = stream.tell()
+            stream.seek(0, io.SEEK_END)
+            end = stream.tell()
+            stream.seek(pos, io.SEEK_SET)
+        except (OSError, AttributeError, io.UnsupportedOperation):
+            return None
+        return max(0, end - pos)
+
     @classmethod
     def shift_chunks(
         cls, chunks: T.Iterable[bytes], offset: int

tests/unit/test_upload_api_v4.py

@@ -3,6 +3,8 @@
 # This source code is licensed under the BSD license found in the
 # LICENSE file in the root directory of this source tree.
 
+from __future__ import annotations
+
 import io
 from pathlib import Path
 
@@ -69,3 +71,24 @@ def _gen_chunks():
     # reupload should not affect the file
     upload_service.upload_chunks(_gen_chunks())
     assert (tmpdir.join("FOOBAR2.txt").read_binary()) == b"foobar"
+
+
+class _ReadSpy(io.BytesIO):
+    def __init__(self, data: bytes):
+        super().__init__(data)
+        self.read_sizes: list[int] = []
+
+    def read(self, size: int | None = -1) -> bytes:  # type: ignore[override]
+        if size is None:
+            size = -1
+        self.read_sizes.append(size)
+        return super().read(size)
+
+
+def test_chunkize_caps_read_to_remaining_bytes():
+    spy = _ReadSpy(b"hello world")
+    chunks = list(
+        upload_api_v4.UploadService.chunkize_byte_stream(spy, 1024 * 1024 * 1000)
+    )
+    assert b"".join(chunks) == b"hello world"
+    assert spy.read_sizes == [len(b"hello world")]