From 13b521651bdbe31bf0179898646c04391d46f20b Mon Sep 17 00:00:00 2001 From: Ruben Fiszel Date: Mon, 10 Aug 2026 23:04:27 +0200 Subject: [PATCH] fix(python-client): return at most size bytes from S3BufferedReader.read (#10623) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * feat: add unit tests for S3BufferedReader.read and improve read method implementation * feat: refactor S3BufferedReader.read method and add unit tests for its functionality * feat: implement peek() on S3BufferedReader with buffered reads * fix(python-client): keep the read(size) contract and trim the test surface Drop the duplicated `TestS3BufferedReaderRead` class from `python-client/tests/wmill_client_test.py`: CI runs `pytest tests/` from `python-client/wmill`, so that legacy manual harness never executes, and the same assertions already live in `python-client/wmill/tests/test_s3_reader.py`. Narrow that file to the four behaviours a future change could break, and make the `bytes_generator` guard actually call `bytes_generator`. Align `peek()` with `io.BufferedReader.peek`, which does at most one read on the underlying stream, rather than looping until `size` bytes are buffered. Co-Authored-By: Claude Opus 5 (1M context) * fix(python-client): hold read1 to one underlying read read1 forwarded to read, so read1(-1) drained the whole object — the same unbounded buffering this branch removes from read. Now that a buffer exists, read1 can honour its own contract: fill only when the buffer is empty, then serve from it. Also treat read(None) as read(-1), per the BufferedReader contract, and pin that read(0) does not pull from the stream: that holds only because the drain sentinel is a negative size, and widening it to any falsy size would reintroduce whole-file buffering. Co-Authored-By: Claude Opus 5 (1M context) * fix(python-client): return from read1(0) without touching the stream A zero-length read has nothing to serve, so pulling a chunk to satisfy it both wastes a round trip and advances the stream. Guard it ahead of the fill, and pin it with a chunk source that counts pulls. Co-Authored-By: Claude Opus 5 (1M context) --------- Co-authored-by: Tushar Co-authored-by: Claude Opus 5 (1M context) --- python-client/wmill/tests/test_s3_reader.py | 97 +++++++++++++++++++++ python-client/wmill/wmill/s3_reader.py | 57 +++++++++--- 2 files changed, 140 insertions(+), 14 deletions(-) create mode 100644 python-client/wmill/tests/test_s3_reader.py diff --git a/python-client/wmill/tests/test_s3_reader.py b/python-client/wmill/tests/test_s3_reader.py new file mode 100644 index 0000000000..0314791244 --- /dev/null +++ b/python-client/wmill/tests/test_s3_reader.py @@ -0,0 +1,97 @@ +"""Unit tests for S3BufferedReader: no network or env needed.""" + +from wmill.s3_reader import S3BufferedReader, bytes_generator + +CHUNKS = [b"AAAAAAAAAA", b"BBBBBBBBBB", b"CCCCCCCCCC"] + + +class _FakeStream: + """Stands in for the httpx streaming response the reader consumes.""" + + status_code = 200 + + def __init__(self, chunks): + self._chunks = chunks + + def __enter__(self): + return self + + def __exit__(self, *args): + pass + + def iter_bytes(self): + return iter(self._chunks) + + +class CountingIterator: + """Chunk source that records how many times the reader pulled from it.""" + + def __init__(self, chunks): + self._chunks = chunks + self.pulls = 0 + + def __iter__(self): + for chunk in self._chunks: + self.pulls += 1 + yield chunk + + +class _FakeClient: + """Stands in for the httpx client, so construction never touches the network.""" + + def __init__(self, chunks): + self._chunks = chunks + + def stream(self, method, url, params=None, timeout=None): + return _FakeStream(self._chunks) + + +def make_reader(chunks): + reader = S3BufferedReader("ws", _FakeClient(chunks), "file.txt", None, None) + reader.__enter__() + return reader + + +def test_read_size_slices_chunks_and_keeps_the_remainder(): + reader = make_reader(CHUNKS) + # read(0) must not pull from the stream: the "drain everything" sentinel is + # a negative size, and widening it to any falsy size would reintroduce the + # whole-file buffering this reader is built to avoid. + assert reader.read(0) == b"" + assert reader.read(5) == b"AAAAA" + assert reader.read(5) == b"AAAAA" + assert reader.read(10) == b"BBBBBBBBBB" + assert reader.read(7) == b"CCCCCCC" + assert reader.read(5) == b"CCC" + assert reader.read(5) == b"" + + +def test_read_all_drains_both_the_buffer_and_the_stream(): + reader = make_reader(CHUNKS) + assert reader.read(5) == b"AAAAA" + assert reader.read(-1) == b"AAAAABBBBBBBBBBCCCCCCCCCC" + + +def test_bytes_generator_yields_50kb_slices_of_64kb_chunks(): + reader = make_reader([b"x" * 65536] * 5) + sizes = [len(chunk) for chunk in bytes_generator(reader)] + assert max(sizes) <= 50 * 1024 + assert sum(sizes) == 5 * 65536 + + +def test_peek_does_not_consume(): + reader = make_reader(CHUNKS) + assert reader.peek() == b"AAAAAAAAAA" + assert reader.read(10) == b"AAAAAAAAAA" + + +def test_read1_stops_after_one_chunk(): + counting = CountingIterator(CHUNKS) + reader = make_reader(counting) + # A zero-length read must not touch the stream at all. + assert reader.read1(0) == b"" + assert counting.pulls == 0 + # read1(-1) must not drain the stream the way read(-1) does. + assert reader.read1(-1) == b"AAAAAAAAAA" + assert reader.read1(4) == b"BBBB" + assert counting.pulls == 2 diff --git a/python-client/wmill/wmill/s3_reader.py b/python-client/wmill/wmill/s3_reader.py index 1f3616c746..2a8f3de2c8 100644 --- a/python-client/wmill/wmill/s3_reader.py +++ b/python-client/wmill/wmill/s3_reader.py @@ -28,6 +28,7 @@ class S3BufferedReader(BufferedReader): params=params, timeout=None, ) + self._buffer = bytearray() def __enter__(self): reader = self._context_manager.__enter__() @@ -46,25 +47,53 @@ class S3BufferedReader(BufferedReader): return self def peek(self, size=0): - raise Exception("Not implemented, use read() instead") + """Return buffered bytes without consuming them. + + Reads the underlying stream at most once, so the amount returned may be + more or less than `size`. + """ + if not self._buffer: + self._fill(1) + return bytes(self._buffer) + + def _fill(self, limit): + # iter_bytes() yields whole HTTP chunks (~64KB), so a caller asking for + # `limit` bytes has to accumulate until the buffer holds that many. + # A negative limit means drain the stream. + while limit < 0 or len(self._buffer) < limit: + try: + self._buffer.extend(next(self._iterator)) + except StopIteration: + break def read(self, size=-1): - read_result = [] + # BufferedReader.read(None) is documented as equivalent to read(-1). + if size is None: + size = -1 + self._fill(size) if size < 0: - for b in self._iterator: - read_result.append(b) - else: - for i in range(size): - try: - b = self._iterator.__next__() - except StopIteration: - break - read_result.append(b) - - return b"".join(read_result) + result = bytes(self._buffer) + self._buffer.clear() + return result + result = bytes(self._buffer[:size]) + del self._buffer[:size] + return result def read1(self, size=-1): - return self.read(size) + """Return up to `size` bytes, reading the underlying stream at most once. + + Unlike `read`, a negative `size` returns only what is already buffered + rather than draining the whole object. + """ + if size == 0: + return b"" + if not self._buffer: + self._fill(1) + if size is None or size < 0: + size = len(self._buffer) + result = bytes(self._buffer[:size]) + del self._buffer[:size] + return result def __exit__(self, *args): self._context_manager.__exit__(*args)