diff --git a/mapillary_tools/upload_api_v4.py b/mapillary_tools/upload_api_v4.py index 569c5055..5f8f7c42 100644 --- a/mapillary_tools/upload_api_v4.py +++ b/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 diff --git a/tests/unit/test_upload_api_v4.py b/tests/unit/test_upload_api_v4.py index da4847ee..b6473f97 100644 --- a/tests/unit/test_upload_api_v4.py +++ b/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")]