fix(python-client): return at most size bytes from S3BufferedReader.read (#10623)

* 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) <noreply@anthropic.com>

* 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) <noreply@anthropic.com>

* 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) <noreply@anthropic.com>

---------

Co-authored-by: Tushar <tusharanshu18@gmail.com>
Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com>

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

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)