Catch asyncio.CancelledError and raise GeneratorExit in aiter functions

Open
  • Is something wrong?
#797 0 comments 0 reactions 0 assignees View on GitHub

Maintainers usually reply within 1 day

Nobody has claimed this yet.

Assessment

Difficulty
3/5
Estimated time
1-2 days
Newbie friendliness
65/100
Issue type
Bug
Clarity
Mostly clear
Activity status
Quiet
Tech stack
python
Domain
api

Research direction

Start in httpx/_client.py by reading the aiter_raw and aiter_bytes functions and how their async streams are closed. Reproduce cancellation with generator.close() or asyncio.gather(..., return_exceptions=True), then verify that cancellation produces GeneratorExit while response cleanup still occurs.

Written by the indexing model from the issue text.

Description

Originally opened by @gjmhmm8 on 2023-11-24 18:32:59 in encode/httpx

Hello, I am using httpx to make asynchronous requests in a generator function. However, I encountered a problem when I tried to cancel the generator using generator.close() or asyncio.gather(*tasks, return_exceptions=True). The generator did not exit gracefully and raised an asyncio.CancelledError instead of a GeneratorExit. This caused some unwanted side effects and made it difficult to handle the cancellation properly.

I looked into the source code of httpx and found that the problem was in the aiter_raw and aiter_bytes functions in the httpx/_client.py module. These functions use async for to iterate over the response stream, but they do not catch the asyncio.CancelledError that may be raised when the stream is cancelled. According to the [documentation] of asyncio.CancelledError, this exception should be caught and either re-raised or suppressed. In this case, I think it would make sense to re-raise it as a GeneratorExit, which is the expected exception for generator termination.

I suggest adding a try-except block around the async for loop in the aiter_raw and aiter_bytes functions, like this:
async def aiter_raw(
self, chunk_size: typing.Optional[int] = None
) -> typing.AsyncIterator[bytes]:
"""
A byte-iterator over the raw response content.
"""
if self.is_stream_consumed:
raise StreamConsumed()
if self.is_closed:
raise StreamClosed()
if not isinstance(self.stream, AsyncByteStream):
raise RuntimeError("Attempted to call an async iterator on an sync stream.")

    self.is_stream_consumed = True
    self._num_bytes_downloaded = 0
    chunker = ByteChunker(chunk_size=chunk_size)
    try:
        with request_context(request=self._request):
            async for raw_stream_bytes in self.stream:
                self._num_bytes_downloaded += len(raw_stream_bytes)
                for chunk in chunker.decode(raw_stream_bytes):
                    await asyncio.sleep(0)
                    yield chunk
        for chunk in chunker.flush():
            await asyncio.sleep(0)
            yield chunk
    except asyncio.CancelledError:
        raise GeneratorExit
    finally:
        await self.aclose()
Dominant language
Python
Stars
1.6k
Forks
103
Avg merge
1d 19h
Merged PRs (30d)
25

Getting set up

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

More from pydantic/httpx2

All issues in pydantic/httpx2

Similar issues

More Python issues

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.