add support for streaming with TaskGroups

[graingert]

Code of Conduct

• I agree to follow Django's Code of Conduct

Feature Description

see https://forum.djangoproject.com/t/streamingresponse-driven-by-a-taskgroup/40320/4

I'd like to be able to write code that combines multiple streams of data:

async def news_and_weather(request: HttpRequest) -> StreamingHttpResponse:
    async def gen() -> AsyncGenerator[bytes]:
        async def push(ws_url: str, tx: MemoryObjectSendStream) -> None:
            async with tx, connect_ws(ws_url) as conn:
                async for msg in conn:
                    await tx.send(msg)

        async with anyio.create_task_group() as tg:
            tx, rx =  anyio.create_memory_object_stream[bytes]()
            with tx, rx:
                tg.start_soon(push, "ws://example.com/news", tx.clone())
                tg.start_soon(push, "ws://example.com/weather", tx.clone())
                tx.close()
                async for msg in rx:
                    yield msg  # yield in async generator!! illegal inside TaskGroup!
    return StreamingHttpResponse(gen())

Problem

however this doesn’t work because I’m using a yield inside an async generator that’s not a context manager, and calling aclosing() on that async generator is not sufficient to allow a TaskGroup to cancel itself and catch the cancel error.

from useful_types import SupportsAnext

class AsyncIteratorBytesResource(Protocol):
    """
    all the machinery needed to safely run an AsyncGenerator[Bytes]

    (for django-stubs) this allows AsyncGenerator[bytes] but is less strict
    so would also allow a anyio MemoryObjectRecieveStream[bytes]]
    """

    async def __aiter__(self) -> SupportsAnext[bytes]: ...
    async def aclose(self) -> object: ...


async def news_and_weather(request: HttpRequest) -> StreamingAcmgrHttpResponse:
    @contextlib.asynccontextmanager
    async def acmgr_gen() -> AsyncGenerator[AsyncIteratorBytesResource]:
        async def push(ws_url: str, tx: MemoryObjectSendStream) -> None:
            async with tx, connect_ws(ws_url) as conn:
                async for msg in conn:
                    await tx.send(msg)

        async with anyio.create_task_group() as tg:
            tx, rx =  anyio.create_memory_object_stream[bytes]()
            with tx, rx:
                tg.start_soon(push, "ws://example.com/news", tx.clone())
                tg.start_soon(push, "ws://example.com/weather", tx.clone())
                tx.close()
                yield rx  # yield inside asynccontextmanager, permitted inside TaskGroup

    return StreamingAcmgrHttpResponse(acmgr_gen())

Request or proposal

proposal

Additional Details

No response

Implementation Suggestions

https://github.com/django/django/pull/19364/changes

[github-actions]

Thank you graingert for sharing your idea! We have a lot of them so please be patient. You can see the current queue here.

Community instructions

For commenters, please use the emoji reactions on the issue to express support, and/or concern easily. Please use the comments to ask questions or contribute knowledge about the idea. It is unhelpful to post comments of "I'd love this" or "What's the state of this?"

Reaction Guide

• 👍 This is something I would use
• 👎 This is something that would cause problems for me or Django
• 😕 I’m indifferent to this
• 🎉 This is an easy win

[graingert]

I updated the type annotations to include a new Protocol AsyncIteratorBytesResource for the object that streaming acmgrs must manage

[carltongibson]

This would be a good candidate for allowing type hints where appropriate #23 (and a maybe GSoC idea).

I did the 👍 But this is a great idea.

[graingert]

Here's some background on why this feature is needed https://peps.python.org/pep-0789/

[graingert]

Here's an example with asyncio primitives instead of anyio:

class Eof(object):
    pass

async def news_and_weather(request: HttpRequest) -> StreamingAcmgrHttpResponse:
    @contextlib.asynccontextmanager
    async def acmgr_gen() -> AsyncGenerator[AsyncGenerator[bytes]]:
        async def push(ws_url: str, q: asyncio.Queue[Eof|bytes], eof: Eof) -> None:
            async with connect_ws(ws_url) as conn:
                async for msg in conn:
                    await q.put(msg)
                await q.put(eof)

        async with asyncio.TaskGroup() as tg:
            q = asyncio.Queue[Eof|bytes](0)
            eof_1 = Eof()
            eof_2 = Eof()
            eofs = {eof_1, eof_2}
            
            tg.create_task(push("ws://example.com/news", q, eof_1))
            tg.create_task(push("ws://example.com/weather", q, eof_2))

            async def agen_fn() -> AsyncGenerator[bytes]:
                while True:
                    v = await q.get()
                    if isinstance(v, bytes):
                        yield bytes
                    else:
                        eofs.discard(v)
                        if not eofs:
                            return
                
            
            yield agen_fn()  # yield inside asynccontextmanager, permitted inside TaskGroup

    return StreamingAcmgrHttpResponse(acmgr_gen())

[hartungstenio]

Out of curiosity: why calling aclose on the async generator is not sufficient? It should raise a GeneratorExit inside the TaskGroup, which could be handled.

StreamingHttpResponse seem not to close async generators, only regular generators (by adding generator.close to response._resource_closers.

[Arfey]

Out of curiosity: why calling aclose on the async generator is not sufficient? It should raise a GeneratorExit inside the TaskGroup, which could be handled.

async def generator():
    try:
        for i in range(10):
            yield
    except BaseException as e:
        print(type(e))
        raise e

async def main():
    async with aclosing(generator()) as gen:
        async for _ in gen:
            break

asyncio.run(main())

output

<class 'GeneratorExit'>

async def main():
    async def generator():
        try:
            async with asyncio.TaskGroup():
                yield
        except BaseException as e:
            print(type(e))
            raise e
            
    async with aclosing(generator()) as gen:
        async for _ in gen:
            break

asyncio.run(main())

output

<class 'BaseExceptionGroup'>



  + Exception Group Traceback (most recent call last):
  |   File "/opt/homebrew/Cellar/ipython/8.30.0/libexec/lib/python3.13/site-packages/IPython/core/interactiveshell.py", line 3577, in run_code
  |     exec(code_obj, self.user_global_ns, self.user_ns)
  |     ~~~~^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  |   File "<ipython-input-70-dd7ee55f92ca>", line 14, in <module>
  |     asyncio.run(main())

Each exception raised inside a TaskGroup block is wrapped into a BaseExceptionGroup, while the generator expects a GeneratorExit. Because of this mismatch, it cannot finish properly.

but we can try something like that

async def background():
    try:
        await asyncio.sleep(1)
    except asyncio.CancelledError:
        print('canceled success')


async def main():
    async def generator():
        try:
            async with asyncio.TaskGroup() as tg:
                tg.create_task(background())
                
                # just to release event loop and start background task
                await asyncio.sleep(0)

                for i in range(10):
                    yield i

        except BaseExceptionGroup as eg:
            gen_exits, others = eg.split(GeneratorExit)

            if gen_exits and not others and len(gen_exits.exceptions) == 1:
                raise gen_exits.exceptions[0]
            
            raise
            
    async with aclosing(generator()) as gen:
        async for _ in gen:
            # raise GeneratorExit error
            break

asyncio.run(main())

You are manually unwrapping the GeneratorExit signal from the TaskGroup error box to let the generator shut down silently instead of crashing.

[carltongibson]

Updated to In Progress to match PR and ticket state.