Skip to content

Commit

Permalink
When chunk_size=None, Response.(a)iter_bytes returns data as it arriv…
Browse files Browse the repository at this point in the history
…es in whatever size the chunks are received (encode#394)
  • Loading branch information
cdeler committed Sep 10, 2020
1 parent eb3ea9e commit 60ea230
Show file tree
Hide file tree
Showing 2 changed files with 22 additions and 8 deletions.
20 changes: 16 additions & 4 deletions httpx/_models.py
Original file line number Diff line number Diff line change
Expand Up @@ -915,12 +915,17 @@ def read(self) -> bytes:
self._content = b"".join(self.iter_bytes())
return self._content

def iter_bytes(self, chunk_size: int = 512) -> typing.Iterator[bytes]:
def iter_bytes(
self, chunk_size: typing.Optional[int] = 512
) -> typing.Iterator[bytes]:
"""
A byte-iterator over the decoded response content.
This allows us to handle gzip, deflate, and brotli encoded responses.
"""
yield from drain_by_chunks(self.__iter_bytes(), chunk_size)
if chunk_size is None:
yield from self.__iter_bytes()
else:
yield from drain_by_chunks(self.__iter_bytes(), chunk_size)

def __iter_bytes(self) -> typing.Iterator[bytes]:
if hasattr(self, "_content"):
Expand Down Expand Up @@ -1001,12 +1006,19 @@ 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, chunk_size: int = 512) -> typing.AsyncIterator[bytes]:
async def aiter_bytes(
self, chunk_size: typing.Optional[int] = 512
) -> typing.AsyncIterator[bytes]:
"""
A byte-iterator over the decoded response content.
This allows us to handle gzip, deflate, and brotli encoded responses.
"""
async for chunk in async_drain_by_chunks(self.__aiter_bytes(), chunk_size):
if chunk_size is None:
aiterator = self.__aiter_bytes()
else:
aiterator = async_drain_by_chunks(self.__aiter_bytes(), chunk_size)

async for chunk in aiterator:
yield chunk

async def __aiter_bytes(self) -> typing.AsyncIterator[bytes]:
Expand Down
10 changes: 6 additions & 4 deletions tests/models/test_responses.py
Original file line number Diff line number Diff line change
Expand Up @@ -270,27 +270,29 @@ async def test_aiter_raw_increments_updates_counter():
num_downloaded = response.num_bytes_downloaded


def test_iter_bytes():
@pytest.mark.parametrize("chunk_size", [2, 20, None])
def test_iter_bytes(chunk_size):
response = httpx.Response(
200,
content=b"Hello, world!",
)

content = b""
for part in response.iter_bytes():
for part in response.iter_bytes(chunk_size):
content += part
assert content == b"Hello, world!"


@pytest.mark.asyncio
async def test_aiter_bytes():
@pytest.mark.parametrize("chunk_size", [2, 20, None])
async def test_aiter_bytes(chunk_size):
response = httpx.Response(
200,
content=b"Hello, world!",
)

content = b""
async for part in response.aiter_bytes():
async for part in response.aiter_bytes(chunk_size):
content += part
assert content == b"Hello, world!"

Expand Down

0 comments on commit 60ea230

Please sign in to comment.