Repository: encode/httpx
Pull request: Support for `chunk_size`
Closes https://github.com/encode/httpx/issues/394

- [x] Add `iter_raw(chunk_size: int=None)`
- [x] Add `iter_bytes(chunk_size: int=None)`
- [x] Add `iter_text(chunk_size: int=None)`

First bit of work towards `chunk_size` on the `response.iter_[raw|bytes|text]` methods,
using a decoder class in line with the other decoding.

Nice thing about this is that it's really easy to unit test, eg...

```python
def test_byte_chunker():
    decoder = ByteChunker()
    assert decoder.decode(b"1234567") == [b"1234567"]
    assert decoder.decode(b"89") == [b"89"]
    assert decoder.flush() == []

    decoder = ByteChunker(chunk_size=3)
    assert decoder.decode(b"1234567") == [b"123", b"456"]
    assert decoder.decode(b"89") == [b"789"]
    assert decoder.flush() == []

    decoder = ByteChunker(chunk_size=3)
    assert decoder.decode(b"123456") == [b"123", b"456"]
    assert decoder.decode(b"789") == [b"789"]
    assert decoder.flush() == []

    decoder = ByteChunker(chunk_size=3)
    assert decoder.decode(b"123456") == [b"123", b"456"]
    assert decoder.decode(b"78") == []
    assert decoder.flush() == [b"78"]
```

Problem statement:
Allow setting `chunk_size` for `Response.iter_bytes()` etc...
Requests allowed setting chunk_size within `.iter_content()` which is currently not an option for our alternatives `.stream()` and `.stream_text()`.

For `.stream_text()` we should go the extra step and fix the issue that users sometimes run into when using this feature and use chunk-size for measuring the decoded text, not the raw bytes.

Hints from issue discussion:
I seem to remeber this playing into some primitives around streaming bytes vs text that we never ended up digging into?
A good first pass onto this would be to change the decoder interface slightly, so that instead of eg. yeilding a byte chunk, they yield a *list* of byte chunks.

On the first refactoring pass, we don't need to *actually* change the internal implmentation much - the decoders can just always yield a list with a single item.

We'd *then* be able to add a `chunk_size` argument to the decoders, which would return 0, 1, or many properly-sized chunks on each yield.
Updated the issue title to reflect the current `Response.aiter_*` API :-) (see #610).
How could I help with this issue? 
Hi @b0g3r! I think this is still something we’d like to have, and given discussions in https://github.com/python-gitlab/python-gitlab/pull/1036 it seems like some folks would like to see it too. :)

Ways to move forward would be:

- Propose an API for this, with context on the existing API on Requests
- Investigate implementation details (ie how are we going to split chunks: buffering, other? Looking at how Requests/urllib3 do this can help)
- Draft a PR :)
Do I understand correctly that we will need to forward chunk_size here?
https://github.com/encode/httpx/blob/a82adcc933345c6b8cb1623b031eb85723e7665b/httpx/_dispatch/urllib3.py#L112-L115
@b0g3r Careful that we're in a sort of transition state w.r.t. urllib3 usage due to #804 (we'll soon use our own sync implementation, though keeping urllib3 as an option). Due to this I wouldn't advise relying on any existing urllib3 functionality — also because we'd want to provide chunk sizing on the async layer too, and it'd be odd to have a different implementation on both sides.

I think we want to look at controlling the chunk size directly from `response.iter_bytes()`/`response.aiter_bytes()`, instead…
@b0g3r So, as with comment https://github.com/encode/httpx/issues/394#issuecomment-567899958 - the right place to start with this would be a pull request to https://github.com/encode/httpx/blob/master/httpx/_decoders.py that changes the interface of the decoders, so that they return a *list of bytes* rather than *bytes*.

(And correspondingly, changing the places where the response calls the decoder such as https://github.com/encode/httpx/blob/a82adcc933345c6b8cb1623b031eb85723e7665b/httpx/_models.py#L915 to deal with a list of bytes as a return result.)

I'd start with that as a foundational pull request, which will then make the remaining work much easier. (Adding chunk sizes to the decoder interface, and through to the response methods.)
### (a)iter_raw(self, chunk_size=1)
- `chunk_size=1` because [requests.Response.iter_content](https://2.python-requests.org/en/master/api/#requests.Response.iter_content) has it
- instead of
```python 
for part in self._raw_stream:
    yield part
```
let's use bytestring as buffer
```python
buffer = b""
for part in self._raw_stream:
    buffer += part
    while len(buffer) >= chunk_size:
        yield buffer[:chunk_size]
        buffer = buffer[chunk_size:]
if buffer:
    yield buffer
```

### (a)iter_bytes(self, chunk_size=ITER_CHUNK_SIZE)
 - `chunk_size=ITER_CHUNK_SIZE` (512) because [requests](https://2.python-requests.org/en/master/api/#requests.Response.iter_lines) has it 🌚 
 - calls `(a)iter_raw `

### (a)iter_text(self, chunk_size=ITER_CHUNK_SIZE)
 - calls `(a)iter_bytes`

### (a)iter_line(self, chunk_size=ITER_CHUNK_SIZE)
 - calls `(a)iter_test`
 - current code expects that each chunk containt full line(s), but it's not true (UPD: found code for splitting in LineDecoder)
 - requests has elegant [solution](https://github.com/psf/requests/blob/2758124a13cddff7244b97b5ffe3fddabb90bc18/requests/models.py#L793-L814)
@tomchristie As I see `(a)iter_raw` doesn't use any decoder 🤔

Commit message: Add TextChunker

Patch to review:
diff --git a/httpx/_decoders.py b/httpx/_decoders.py
index 40c6da9fc0..c5b3764b88 100644
--- a/httpx/_decoders.py
+++ b/httpx/_decoders.py
@@ -4,6 +4,7 @@
 See: https://developer.mozilla.org/en-US/docs/Web/HTTP/Headers/Content-Encoding
 """
 import codecs
+import io
 import typing
 import zlib
 
@@ -157,6 +158,84 @@ def flush(self) -> bytes:
         return data
 
 
+class ByteChunker:
+    """
+    Handles returning byte content in fixed-size chunks.
+    """
+
+    def __init__(self, chunk_size: int = None) -> None:
+        self._buffer = io.BytesIO()
+        self._chunk_size = chunk_size
+
+    def decode(self, content: bytes) -> typing.List[bytes]:
+        if self._chunk_size is None:
+            return [content]
+
+        self._buffer.write(content)
+        if self._buffer.tell() >= self._chunk_size:
+            value = self._buffer.getvalue()
+            chunks = [
+                value[i : i + self._chunk_size]
+                for i in range(0, len(value), self._chunk_size)
+            ]
+            if len(chunks[-1]) == self._chunk_size:
+                self._buffer.seek(0)
+                self._buffer.truncate()
+                return chunks
+            else:
+                self._buffer.seek(0)
+                self._buffer.write(chunks[-1])
+                self._buffer.truncate()
+                return chunks[:-1]
+        else:
+            return []
+
+    def flush(self) -> typing.List[bytes]:
+        value = self._buffer.getvalue()
+        self._buffer.seek(0)
+        self._buffer.truncate()
+        return [value] if value else []
+
+
+class TextChunker:
+    """
+    Handles returning text content in fixed-size chunks.
+    """
+
+    def __init__(self, chunk_size: int = None) -> None:
+        self._buffer = io.StringIO()
+        self._chunk_size = chunk_size
+
+    def decode(self, content: str) -> typing.List[str]:
+        if self._chunk_size is None:
+            return [content]
+
+        self._buffer.write(content)
+        if self._buffer.tell() >= self._chunk_size:
+            value = self._buffer.getvalue()
+            chunks = [
+                value[i : i + self._chunk_size]
+                for i in range(0, len(value), self._chunk_size)
+            ]
+            if len(chunks[-1]) == self._chunk_size:
+                self._buffer.seek(0)
+                self._buffer.truncate()
+                return chunks
+            else:
+                self._buffer.seek(0)
+                self._buffer.write(chunks[-1])
+                self._buffer.truncate()
+                return chunks[:-1]
+        else:
+            return []
+
+    def flush(self) -> typing.List[str]:
+        value = self._buffer.getvalue()
+        self._buffer.seek(0)
+        self._buffer.truncate()
+        return [value] if value else []
+
+
 class TextDecoder:
     """
     Handles incrementally decoding bytes into text
diff --git a/httpx/_models.py b/httpx/_models.py
index 6d90f90b4e..a51590787c 100644
--- a/httpx/_models.py
+++ b/httpx/_models.py
@@ -17,6 +17,7 @@
 from ._content_streams import ByteStream, ContentStream, encode
 from ._decoders import (
     SUPPORTED_DECODERS,
+    ByteChunker,
     ContentDecoder,
     IdentityDecoder,
     LineDecoder,
@@ -912,19 +913,28 @@ def read(self) -> bytes:
             self._content = b"".join(self.iter_bytes())
         return self._content
 
-    def iter_bytes(self) -> typing.Iterator[bytes]:
+    def iter_bytes(self, chunk_size: int = None) -> typing.Iterator[bytes]:
         """
         A byte-iterator over the decoded response content.
         This allows us to handle gzip, deflate, and brotli encoded responses.
         """
         if hasattr(self, "_content"):
-            yield self._content
+            chunk_size = len(self._content) if chunk_size is None else chunk_size
+            for i in range(0, len(self._content), chunk_size):
+                yield self._content[i : i + chunk_size]
         else:
             decoder = self._get_content_decoder()
+            chunker = ByteChunker(chunk_size=chunk_size)
             with self._wrap_decoder_errors():
-                for chunk in self.iter_raw():
-                    yield decoder.decode(chunk)
-                yield decoder.flush()
+                for raw_bytes in self.iter_raw():
+                    decoded = decoder.decode(raw_bytes)
+                    for chunk in chunker.decode(decoded):
+                        yield chunk
+                decoded = decoder.flush()
+                for chunk in chunker.decode(decoded):
+                    yield chunk
+                for chunk in chunker.flush():
+                    yield chunk
 
     def iter_text(self) -> typing.Iterator[str]:
         """
@@ -947,7 +957,7 @@ def iter_lines(self) -> typing.Iterator[str]:
             for line in decoder.flush():
                 yield line
 
-    def iter_raw(self) -> typing.Iterator[bytes]:
+    def iter_raw(self, chunk_size: int = None) -> typing.Iterator[bytes]:
         """
         A byte-iterator over the raw response content.
         """
@@ -958,10 +968,17 @@ def iter_raw(self) -> typing.Iterator[bytes]:
 
         self.is_stream_consumed = True
         self._num_bytes_downloaded = 0
+        chunker = ByteChunker(chunk_size=chunk_size)
+
         with map_exceptions(HTTPCORE_EXC_MAP, request=self._request):
-            for part in self._raw_stream:
-                self._num_bytes_downloaded += len(part)
-                yield part
+            for raw_stream_bytes in self._raw_stream:
+                self._num_bytes_downloaded += len(raw_stream_bytes)
+                for chunk in chunker.decode(raw_stream_bytes):
+                    yield chunk
+
+        for chunk in chunker.flush():
+            yield chunk
+
         self.close()
 
     def next(self) -> "Response":
@@ -996,19 +1013,28 @@ async def aread(self) -> bytes:
             self._content = b"".join([part async for part in self.aiter_bytes()])
         return self._content
 
-    async def aiter_bytes(self) -> typing.AsyncIterator[bytes]:
+    async def aiter_bytes(self, chunk_size: int = None) -> typing.AsyncIterator[bytes]:
         """
         A byte-iterator over the decoded response content.
         This allows us to handle gzip, deflate, and brotli encoded responses.
         """
         if hasattr(self, "_content"):
-            yield self._content
+            chunk_size = len(self._content) if chunk_size is None else chunk_size
+            for i in range(0, len(self._content), chunk_size):
+                yield self._content[i : i + chunk_size]
         else:
             decoder = self._get_content_decoder()
+            chunker = ByteChunker(chunk_size=chunk_size)
             with self._wrap_decoder_errors():
-                async for chunk in self.aiter_raw():
-                    yield decoder.decode(chunk)
-                yield decoder.flush()
+                async for raw_bytes in self.aiter_raw():
+                    decoded = decoder.decode(raw_bytes)
+                    for chunk in chunker.decode(decoded):
+                        yield chunk
+                decoded = decoder.flush()
+                for chunk in chunker.decode(decoded):
+                    yield chunk
+                for chunk in chunker.flush():
+                    yield chunk
 
     async def aiter_text(self) -> typing.AsyncIterator[str]:
         """
@@ -1031,7 +1057,7 @@ async def aiter_lines(self) -> typing.AsyncIterator[str]:
             for line in decoder.flush():
                 yield line
 
-    async def aiter_raw(self) -> typing.AsyncIterator[bytes]:
+    async def aiter_raw(self, chunk_size: int = None) -> typing.AsyncIterator[bytes]:
         """
         A byte-iterator over the raw response content.
         """
@@ -1042,10 +1068,17 @@ async def aiter_raw(self) -> typing.AsyncIterator[bytes]:
 
         self.is_stream_consumed = True
         self._num_bytes_downloaded = 0
+        chunker = ByteChunker(chunk_size=chunk_size)
+
         with map_exceptions(HTTPCORE_EXC_MAP, request=self._request):
-            async for part in self._raw_stream:
-                self._num_bytes_downloaded += len(part)
-                yield part
+            async for raw_stream_bytes in self._raw_stream:
+                self._num_bytes_downloaded += len(raw_stream_bytes)
+                for chunk in chunker.decode(raw_stream_bytes):
+                    yield chunk
+
+        for chunk in chunker.flush():
+            yield chunk
+
         await self.aclose()
 
     async def anext(self) -> "Response":
diff --git a/tests/models/test_responses.py b/tests/models/test_responses.py
index 30d600086a..f9b4220760 100644
--- a/tests/models/test_responses.py
+++ b/tests/models/test_responses.py
@@ -227,6 +227,26 @@ def test_iter_raw():
     assert raw == b"Hello, world!"
 
 
+def test_iter_raw_with_chunksize():
+    stream = IteratorStream(iterator=streaming_body())
+    response = httpx.Response(200, stream=stream)
+
+    parts = [part for part in response.iter_raw(chunk_size=5)]
+    assert parts == [b"Hello", b", wor", b"ld!"]
+
+    stream = IteratorStream(iterator=streaming_body())
+    response = httpx.Response(200, stream=stream)
+
+    parts = [part for part in response.iter_raw(chunk_size=13)]
+    assert parts == [b"Hello, world!"]
+
+    stream = IteratorStream(iterator=streaming_body())
+    response = httpx.Response(200, stream=stream)
+
+    parts = [part for part in response.iter_raw(chunk_size=20)]
+    assert parts == [b"Hello, world!"]
+
+
 def test_iter_raw_increments_updates_counter():
     stream = IteratorStream(iterator=streaming_body())
 
@@ -255,6 +275,27 @@ async def test_aiter_raw():
     assert raw == b"Hello, world!"
 
 
+@pytest.mark.asyncio
+async def test_aiter_raw_with_chunksize():
+    stream = AsyncIteratorStream(aiterator=async_streaming_body())
+    response = httpx.Response(200, stream=stream)
+
+    parts = [part async for part in response.aiter_raw(chunk_size=5)]
+    assert parts == [b"Hello", b", wor", b"ld!"]
+
+    stream = AsyncIteratorStream(aiterator=async_streaming_body())
+    response = httpx.Response(200, stream=stream)
+
+    parts = [part async for part in response.aiter_raw(chunk_size=13)]
+    assert parts == [b"Hello, world!"]
+
+    stream = AsyncIteratorStream(aiterator=async_streaming_body())
+    response = httpx.Response(200, stream=stream)
+
+    parts = [part async for part in response.aiter_raw(chunk_size=20)]
+    assert parts == [b"Hello, world!"]
+
+
 @pytest.mark.asyncio
 async def test_aiter_raw_increments_updates_counter():
     stream = AsyncIteratorStream(aiterator=async_streaming_body())
@@ -271,10 +312,7 @@ async def test_aiter_raw_increments_updates_counter():
 
 
 def test_iter_bytes():
-    response = httpx.Response(
-        200,
-        content=b"Hello, world!",
-    )
+    response = httpx.Response(200, content=b"Hello, world!")
 
     content = b""
     for part in response.iter_bytes():
@@ -282,6 +320,23 @@ def test_iter_bytes():
     assert content == b"Hello, world!"
 
 
+def test_iter_bytes_with_chunk_size():
+    stream = IteratorStream(iterator=streaming_body())
+    response = httpx.Response(200, stream=stream)
+    parts = [part for part in response.iter_bytes(chunk_size=5)]
+    assert parts == [b"Hello", b", wor", b"ld!"]
+
+    stream = IteratorStream(iterator=streaming_body())
+    response = httpx.Response(200, stream=stream)
+    parts = [part for part in response.iter_bytes(chunk_size=13)]
+    assert parts == [b"Hello, world!"]
+
+    stream = IteratorStream(iterator=streaming_body())
+    response = httpx.Response(200, stream=stream)
+    parts = [part for part in response.iter_bytes(chunk_size=20)]
+    assert parts == [b"Hello, world!"]
+
+
 @pytest.mark.asyncio
 async def test_aiter_bytes():
     response = httpx.Response(
@@ -295,6 +350,24 @@ async def test_aiter_bytes():
     assert content == b"Hello, world!"
 
 
+@pytest.mark.asyncio
+async def test_aiter_bytes_with_chunk_size():
+    stream = AsyncIteratorStream(aiterator=async_streaming_body())
+    response = httpx.Response(200, stream=stream)
+    parts = [part async for part in response.aiter_bytes(chunk_size=5)]
+    assert parts == [b"Hello", b", wor", b"ld!"]
+
+    stream = AsyncIteratorStream(aiterator=async_streaming_body())
+    response = httpx.Response(200, stream=stream)
+    parts = [part async for part in response.aiter_bytes(chunk_size=13)]
+    assert parts == [b"Hello, world!"]
+
+    stream = AsyncIteratorStream(aiterator=async_streaming_body())
+    response = httpx.Response(200, stream=stream)
+    parts = [part async for part in response.aiter_bytes(chunk_size=20)]
+    assert parts == [b"Hello, world!"]
+
+
 def test_iter_text():
     response = httpx.Response(
         200,
diff --git a/tests/test_decoders.py b/tests/test_decoders.py
index dbbaac5450..89fd43633a 100644
--- a/tests/test_decoders.py
+++ b/tests/test_decoders.py
@@ -7,10 +7,12 @@
 from httpx._content_streams import AsyncIteratorStream
 from httpx._decoders import (
     BrotliDecoder,
+    ByteChunker,
     DeflateDecoder,
     GZipDecoder,
     IdentityDecoder,
     LineDecoder,
+    TextChunker,
     TextDecoder,
 )
 
@@ -313,6 +315,50 @@ def test_line_decoder_crnl():
     assert decoder.flush() == []
 
 
+def test_byte_chunker():
+    decoder = ByteChunker()
+    assert decoder.decode(b"1234567") == [b"1234567"]
+    assert decoder.decode(b"89") == [b"89"]
+    assert decoder.flush() == []
+
+    decoder = ByteChunker(chunk_size=3)
+    assert decoder.decode(b"1234567") == [b"123", b"456"]
+    assert decoder.decode(b"89") == [b"789"]
+    assert decoder.flush() == []
+
+    decoder = ByteChunker(chunk_size=3)
+    assert decoder.decode(b"123456") == [b"123", b"456"]
+    assert decoder.decode(b"789") == [b"789"]
+    assert decoder.flush() == []
+
+    decoder = ByteChunker(chunk_size=3)
+    assert decoder.decode(b"123456") == [b"123", b"456"]
+    assert decoder.decode(b"78") == []
+    assert decoder.flush() == [b"78"]
+
+
+def test_text_chunker():
+    decoder = TextChunker()
+    assert decoder.decode("1234567") == ["1234567"]
+    assert decoder.decode("89") == ["89"]
+    assert decoder.flush() == []
+
+    decoder = TextChunker(chunk_size=3)
+    assert decoder.decode("1234567") == ["123", "456"]
+    assert decoder.decode("89") == ["789"]
+    assert decoder.flush() == []
+
+    decoder = TextChunker(chunk_size=3)
+    assert decoder.decode("123456") == ["123", "456"]
+    assert decoder.decode("789") == ["789"]
+    assert decoder.flush() == []
+
+    decoder = TextChunker(chunk_size=3)
+    assert decoder.decode("123456") == ["123", "456"]
+    assert decoder.decode("78") == []
+    assert decoder.flush() == ["78"]
+
+
 def test_invalid_content_encoding_header():
     headers = [(b"Content-Encoding", b"invalid-header")]
     body = b"test 123"


Code review comments:
diff_hunk: @@ -912,19 +913,28 @@ def read(self) -> bytes:
             self._content = b"".join(self.iter_bytes())
         return self._content
 
-    def iter_bytes(self) -> typing.Iterator[bytes]:
+    def iter_bytes(self, chunk_size: int = None) -> typing.Iterator[bytes]:
path: httpx/_models.py
text: @user1:
Are you sure about chunk_size with default `None`?

 I do agree that it looks more suitable, but `requests` provides us with defaults [chunk_size=1](https://2.python-requests.org/en/master/_modules/requests/models/#Response.iter_content) or [512](https://2.python-requests.org/en/master/_modules/requests/models/#Response.iter_lines)

@user2:
In this PR, when `chunk_size=None` we just return the input `content` unchanged, as one single big chunk.

Yes, this would deviate from what Requests seems to do, but:

- Setting a non-`None` default would break backward compatibility on our side.
- Defaulting to "transparently pass the chunk sent by the server" is probably the most reasonable approach _anyway_.

That said, we'd need to add this deviation from Requests to the compatibility guide. :+1: