Add support for streaming with TaskGroups

[Thomas Grainger]

​https://forum.djangoproject.com/t/streamingresponse-driven-by-a-taskgroup/40320/4
​https://github.com/django/new-features/issues/117

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())

Implementation Suggestions

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

[Jacob Walls]

Thanks. Reception on the new-features repo issue seems positive, so though it hasn't moved through any swimlanes there yet, I will speculatively move this one to Accepted assuming it likely will.

[Carlton Gibson]

General approach looks good. I left some comments, and there’s some discussion about a possible alternative approach (but I’m not sure if that’s viable or not, so we should see if that appears.)

[Thomas Grainger]

How about this approach? Working from my phone so posting as a patch rather than pushing directly.

Summary of changes (from Claude)

django/http/response.py

• _set_streaming_content detects async context managers via hasattr(__aenter__, __aexit__), sets self.is_acmgr = True and stores the value in self.__acmgr
• streaming_content property: is_acmgr is an early return at the top (unnested from is_async). Returns an @asynccontextmanager when is_acmgr is True — callers that check is_acmgr use async with response.streaming_content as agen. Both acmgr and regular async paths apply make_bytes and aclose their underlying iterator in a finally block
• StreamingHttpResponse gains __aenter__/__aexit__: for acmgr responses, __aenter__ calls into self.streaming_content (no duplication of awrapper logic needed), stores the context and the yielded generator, then __aexit__ closes the generator and exits the context in the right order so TaskGroup cleans up correctly
• __iter__, __aiter__, and getvalue all raise IsAcmgrException when is_acmgr is True
• IsAcmgrException is exported from django.http

django/core/handlers/asgi.py

• send_response simplified to async with response as content — works for both regular streaming and acmgr responses. aclosing no longer needed

django/middleware/gzip.py

• Adds is_acmgr branch: captures response.streaming_content (an acmgr) and wraps it in a new @asynccontextmanager that feeds the yielded generator into acompress_sequence

django/utils/text.py

• acompress_sequence wraps its entire body in try/finally to aclose its operand on exit if the method is present

django/core/handlers/asgi.py

@@ -4 +4 @@
 import tempfile
 import traceback
 from collections import defaultdict
-from contextlib import aclosing, closing
+from contextlib import closing
 
 from asgiref.sync import ThreadSensitiveContext, sync_to_async
 
@@ -315 +315 @@
         )
         # Streaming responses need to be pinned to their iterator.
         if response.streaming:
-            # - Consume via `__aiter__` and not `streaming_content` directly,
-            #   to allow mapping of a sync iterator.
-            # - Use aclosing() when consuming aiter. See
-            #   https://github.com/python/cpython/commit/6e8dcdaaa49d4313bf9fab9f9923ca5828fbb10e
-            async with aclosing(aiter(response)) as content:
+            async with response as content:
                 async for part in content:
                     for chunk, _ in self.chunk_bytes(part):
                         await send(

django/http/__init__.py

@@ -8 +8 @@
 )
 from django.http.response import (
     BadHeaderError,
+
     FileResponse,
     Http404,
     HttpResponse,
@@ -47 +48 @@
     "HttpResponseServerError",
     "Http404",
     "BadHeaderError",
+
     "JsonResponse",
     "FileResponse",
 ]

django/http/response.py

@@ -7 +7 @@
 import sys
 import time
 import warnings
+
 from email.header import Header
 from http.client import responses
 from urllib.parse import urlsplit
@@ -104 +105 @@
     pass
 
 
+
+
+
+
 class HttpResponseBase:
     """
     An HTTP response base class with dictionary-accessed headers.
@@ -479 +484 @@
 
     @property
     def streaming_content(self):
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
         if self.is_async:
             # pull to lexical scope to capture fixed reference in case
             # streaming_content is set again later.
             _iterator = self._iterator
 
             async def awrapper():
-                async for part in _iterator:
-                    yield self.make_bytes(part)
+                try:
+                    async for part in _iterator:
+                        yield self.make_bytes(part)
+                finally:
+                    if hasattr(_iterator, "aclose"):
+                        await _iterator.aclose()
 
             return awrapper()
         else:
@@ -498 +525 @@
 
     def _set_streaming_content(self, value):
         # Ensure we can never iterate on "value" more than once.
+
+
+
+
+
+
         try:
             self._iterator = iter(value)
             self.is_async = False
@@ -507 +540 @@
         if hasattr(value, "close"):
             self._resource_closers.append(value.close)
 
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
     def __iter__(self):
+
+
+
+
+
         try:
             return iter(self.streaming_content)
         except TypeError:
@@ -528 +582 @@
             return map(self.make_bytes, iter(async_to_sync(to_list)(self._iterator)))
 
     async def __aiter__(self):
+
+
+
+
+
         try:
             async for part in self.streaming_content:
                 yield part
@@ -544 +603 @@
                 yield part
 
     def getvalue(self):
+
+
+
+
+
         return b"".join(self.streaming_content)
 
 

django/middleware/gzip.py

@@ +1 @@
+
+
 from django.utils.cache import patch_vary_headers
 from django.utils.deprecation import MiddlewareMixin
 from django.utils.regex_helper import _lazy_re_compile
@@ -31 +33 @@
             return response
 
         if response.streaming:
-            if response.is_async:
+            if response.is_acmgr:
+                original_acmgr = response.streaming_content
+                max_random_bytes = self.max_random_bytes
+
+                @asynccontextmanager
+                async def compressed_acmgr():
+                    async with original_acmgr as agen:
+                        yield acompress_sequence(
+                            agen,
+                            max_random_bytes=max_random_bytes,
+                        )
+
+                response.streaming_content = compressed_acmgr()
+            elif response.is_async:
                 response.streaming_content = acompress_sequence(
                     response.streaming_content,
                     max_random_bytes=self.max_random_bytes,

django/utils/text.py

@@ -390 +390 @@
 
 
 async def acompress_sequence(sequence, *, max_random_bytes=None):
-    buf = StreamingBuffer()
-    filename = _get_random_filename(max_random_bytes) if max_random_bytes else None
-    with GzipFile(
-        filename=filename, mode="wb", compresslevel=6, fileobj=buf, mtime=0
-    ) as zfile:
-        # Output headers...
+    try:
+        buf = StreamingBuffer()
+        filename = _get_random_filename(max_random_bytes) if max_random_bytes else None
+        with GzipFile(
+            filename=filename, mode="wb", compresslevel=6, fileobj=buf, mtime=0
+        ) as zfile:
+            # Output headers...
+            yield buf.read()
+            async for item in sequence:
+                zfile.write(item)
+                zfile.flush()
+                data = buf.read()
+                if data:
+                    yield data
         yield buf.read()
-        async for item in sequence:
-            zfile.write(item)
-            zfile.flush()
-            data = buf.read()
-            if data:
-                yield data
-    yield buf.read()
+    finally:
+        if hasattr(sequence, "aclose"):
+            await sequence.aclose()

[Thomas Grainger]

Attaching updated patch based on the approach discussed in ​PR #19364.

Summary of changes

django/http/response.py:

• _set_streaming_content detects async context managers via hasattr(__aenter__, __aexit__), sets is_acmgr = True and stores the value
• streaming_content property returns an @asynccontextmanager wrapping make_bytes when is_acmgr is True
• StreamingHttpResponse gains __aenter__/__aexit__: for acmgr responses, enters the streaming_content CM; for regular responses, returns aiter(self). __aexit__ only handles acmgr CM cleanup — aclosing in the ASGI handler handles iterator close
• __iter__, __aiter__, and getvalue raise IsAcmgrException when is_acmgr is True

django/core/handlers/asgi.py:

• Single code path: async with response as agen, aclosing(agen) as content: — works for both regular streaming and acmgr responses

django/middleware/gzip.py:

• Adds is_acmgr branch that wraps the acmgr in a new @asynccontextmanager feeding into acompress_sequence

django/utils/text.py:

• acompress_sequence wraps its body in try/finally to aclose its operand

Tests

• 13 new tests in tests/httpwrappers/tests.py covering: basic acmgr streaming, make_bytes coercion, IsAcmgrException guards on __iter__/__aiter__/getvalue, __aexit__ on error and break, non-acmgr __aenter__ fallback, reassignment, single-producer TaskGroup+Queue, and multi-producer fan-in (news-and-weather pattern)
• 1 new test in tests/middleware/tests.py for gzip compression of acmgr streaming responses

Docs

• New "Streaming with TaskGroup" section in docs/ref/request-response.txt with full news_and_weather example, key points re PEP 789, and anyio note
• Release note in docs/releases/6.1.txt

Patch

django/core/handlers/asgi.py

@@ -318 +318 @@
         )
         # Streaming responses need to be pinned to their iterator.
         if response.streaming:
-            # - Consume via `__aiter__` and not `streaming_content` directly,
-            #   to allow mapping of a sync iterator.
-            # - Use aclosing() when consuming aiter. See
-            #   https://github.com/python/cpython/commit/6e8dcdaaa49d4313bf9fab9f9923ca5828fbb10e
-            async with aclosing(aiter(response)) as content:
+            # Use aclosing() when consuming aiter. See
+            # https://github.com/python/cpython/commit/6e8dcdaaa49d4313bf9fab9f9923ca5828fbb10e
+            async with response as agen, aclosing(agen) as content:
                 async for part in content:
                     for chunk, _ in self.chunk_bytes(part):
                         await send(

django/http/__init__.py

@@ -8 +8 @@
 )
 from django.http.response import (
     BadHeaderError,
+
     FileResponse,
     Http404,
     HttpResponse,
@@ -47 +48 @@
     "HttpResponseServerError",
     "Http404",
     "BadHeaderError",
+
     "JsonResponse",
     "FileResponse",
 ]

django/http/response.py

@@ -7 +7 @@
 import sys
 import time
 import warnings
+
 from email.header import Header
 from http.client import responses
 from urllib.parse import urlsplit
@@ -104 +105 @@
     pass
 
 
+
+
+
+
 class HttpResponseBase:
     """
     An HTTP response base class with dictionary-accessed headers.
@@ -479 +484 @@
 
     @property
     def streaming_content(self):
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
         if self.is_async:
             # pull to lexical scope to capture fixed reference in case
             # streaming_content is set again later.
             _iterator = self._iterator
 
             async def awrapper():
-                async for part in _iterator:
-                    yield self.make_bytes(part)
+                try:
+                    async for part in _iterator:
+                        yield self.make_bytes(part)
+                finally:
+                    if hasattr(_iterator, "aclose"):
+                        await _iterator.aclose()
 
             return awrapper()
         else:
@@ -498 +525 @@
 
     def _set_streaming_content(self, value):
         # Ensure we can never iterate on "value" more than once.
+
+
+
+
+
+
         try:
             self._iterator = iter(value)
             self.is_async = False
@@ -507 +540 @@
         if hasattr(value, "close"):
             self._resource_closers.append(value.close)
 
+
+
+
+
+
+
+
+
+
+
     def __iter__(self):
+
+
+
+
+
         try:
             return iter(self.streaming_content)
         except TypeError:
@@ -528 +576 @@
             return map(self.make_bytes, iter(async_to_sync(to_list)(self._iterator)))
 
     async def __aiter__(self):
+
+
+
+
+
         try:
             async for part in self.streaming_content:
                 yield part
@@ -544 +597 @@
                 yield part
 
     def getvalue(self):
+
+
+
+
+
         return b"".join(self.streaming_content)
 
 

django/middleware/gzip.py

@@ +1 @@
+
+
 from django.utils.cache import patch_vary_headers
 from django.utils.deprecation import MiddlewareMixin
 from django.utils.regex_helper import _lazy_re_compile
@@ -31 +33 @@
             return response
 
         if response.streaming:
-            if response.is_async:
+            if response.is_acmgr:
+                original_acmgr = response.streaming_content
+                max_random_bytes = self.max_random_bytes
+
+                @asynccontextmanager
+                async def compressed_acmgr():
+                    async with original_acmgr as agen:
+                        yield acompress_sequence(
+                            agen,
+                            max_random_bytes=max_random_bytes,
+                        )
+
+                response.streaming_content = compressed_acmgr()
+            elif response.is_async:
                 response.streaming_content = acompress_sequence(
                     response.streaming_content,
                     max_random_bytes=self.max_random_bytes,

django/utils/text.py

@@ -390 +390 @@
 
 
 async def acompress_sequence(sequence, *, max_random_bytes=None):
-    buf = StreamingBuffer()
-    filename = _get_random_filename(max_random_bytes) if max_random_bytes else None
-    with GzipFile(
-        filename=filename, mode="wb", compresslevel=6, fileobj=buf, mtime=0
-    ) as zfile:
-        # Output headers...
+    try:
+        buf = StreamingBuffer()
+        filename = _get_random_filename(max_random_bytes) if max_random_bytes else None
+        with GzipFile(
+            filename=filename, mode="wb", compresslevel=6, fileobj=buf, mtime=0
+        ) as zfile:
+            # Output headers...
+            yield buf.read()
+            async for item in sequence:
+                zfile.write(item)
+                zfile.flush()
+                data = buf.read()
+                if data:
+                    yield data
         yield buf.read()
-        async for item in sequence:
-            zfile.write(item)
-            zfile.flush()
-            data = buf.read()
-            if data:
-                yield data
-    yield buf.read()
+    finally:
+        if hasattr(sequence, "aclose"):
+            await sequence.aclose()
 
 
 # Expression to match some_token and some_token="with spaces" (and similarly

docs/ref/request-response.txt

@@ -1424 +1424 @@
 returning the ``StreamingHttpResponse`` object, then you may also want to
 :ref:`handle disconnections in the view <async-handling-disconnect>` itself.
 
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
 ``FileResponse`` objects
 ========================
 

docs/releases/6.1.txt

@@ -323 +323 @@
 Requests and Responses
 ~~~~~~~~~~~~~~~~~~~~~~
 
+
+
+
+
+
 * :attr:`HttpRequest.multipart_parser_class <django.http.HttpRequest.multipart_parser_class>`
   can now be customized to use a different multipart parser class.
 

tests/asgi/urls.py

@@ -1 +1 @@
 import asyncio
+
 import threading
 import time
 

tests/httpwrappers/tests.py

@@ +1 @@
+
+
 import copy
 import json
 import os
 import pickle
+
 import unittest
 import uuid
 
@@ -16 +19 @@
     HttpResponseNotModified,
     HttpResponsePermanentRedirect,
     HttpResponseRedirect,
+
     JsonResponse,
     QueryDict,
     SimpleCookie,
@@ -808 +812 @@
         with self.assertWarnsMessage(Warning, msg):
             self.assertEqual(b"hello", await anext(aiter(r)))
 
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
     def test_text_attribute_error(self):
         r = StreamingHttpResponse(iter(["hello", "world"]))
         msg = "This %s instance has no `text` attribute." % r.__class__.__name__

tests/middleware/tests.py

@@ +1 @@
+
 import gzip
 import random
 import re
@@ -936 +937 @@
         self.assertEqual(r.get("Content-Encoding"), "gzip")
         self.assertFalse(r.has_header("Content-Length"))
 
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
     def test_compress_streaming_response_unicode(self):
         """
         Compression is performed on responses with streaming Unicode content.