Add support for streaming responses to ASGITransport

httpx/_transports/asgi.py

@@ -4,14 +4,17 @@
 import httpcore
 import sniffio
 
-from .._content_streams import ByteStream
+from .._content_streams import AsyncIteratorStream, ByteStream
 from .._utils import warn_deprecated
 
 if typing.TYPE_CHECKING:  # pragma: no cover
     import asyncio
     import trio
 
     Event = typing.Union[asyncio.Event, trio.Event]
+    BackgroundTask = typing.Union["AsyncioBackgroundTask", "TrioBackgroundTask"]
+
+AsyncCallback = typing.Callable[[], typing.Awaitable]
 
 
 def create_event() -> "Event":
@@ -25,6 +28,91 @@ def create_event() -> "Event":
         return asyncio.Event()
 
 
+class AsyncioBackgroundTask:
+    def __init__(self, async_fn: typing.Callable) -> None:
+        self.async_fn = async_fn
+        self.task: typing.Optional["asyncio.Task"] = None
+
+    async def start(self) -> None:
+        import asyncio
+
+        assert self.task is None
+        self.task = asyncio.create_task(self.async_fn())
+
+    async def aclose(self, exc: Exception = None) -> None:
+        assert self.task is not None
+        if not self.task.done():
+            self.task.cancel()
+        self.task = None
+
+
+class TrioBackgroundTask:
+    def __init__(self, async_fn: typing.Callable) -> None:
+        self.async_fn = async_fn
+        self.nursery_manager: typing.Optional["trio.NurseryManager"] = None
+
+    async def start(self) -> None:
+        import trio
+
+        assert self.nursery_manager is None
+        self.nursery_manager = trio.open_nursery()
+        nursery = await self.nursery_manager.__aenter__()
+        nursery.start_soon(self.async_fn)
+
+    async def aclose(self, exc: Exception = None) -> None:
+        assert self.nursery_manager is not None
+        if exc is not None:
+            await self.nursery_manager.__aexit__(type(exc), exc, exc.__traceback__)
+        else:
+            await self.nursery_manager.__aexit__(None, None, None)
+
+
+def create_background_task(async_fn: typing.Callable) -> "BackgroundTask":
+    if sniffio.current_async_library() == "trio":
+        return TrioBackgroundTask(async_fn)
+    return AsyncioBackgroundTask(async_fn)
+
+
+def create_channel(
+    capacity: int,
+) -> typing.Tuple[
+    typing.Callable[[], typing.Awaitable[bytes]],
+    typing.Callable[[bytes], typing.Awaitable[None]],
+]:
+    if sniffio.current_async_library() == "trio":
+        import trio
+
+        send_channel, receive_channel = trio.open_memory_channel(capacity)
+        return receive_channel.receive, send_channel.send
+    else:
+        import asyncio
+
+        queue: asyncio.Queue[bytes] = asyncio.Queue(capacity)
+        return queue.get, queue.put
+
+
+async def run_until_first_complete(*async_fns: AsyncCallback) -> None:
+    if sniffio.current_async_library() == "trio":
+        import trio
+
+        async with trio.open_nursery() as nursery:
+
+            async def run(async_fn: AsyncCallback) -> None:
+                await async_fn()
+                nursery.cancel_scope.cancel()
+
+            for async_fn in async_fns:
+                nursery.start_soon(run, async_fn)
+
+    else:
+        import asyncio
+
+        coros = {async_fn() for async_fn in async_fns}
+        done, pending = await asyncio.wait(coros, return_when=asyncio.FIRST_COMPLETED)
+        for task in pending:
+            task.cancel()
+
+
 class ASGITransport(httpcore.AsyncHTTPTransport):
     """
     A custom AsyncTransport that handles sending requests directly to an ASGI app.
@@ -95,18 +183,20 @@ async def request(
         }
         status_code = None
         response_headers = None
-        body_parts = []
+        consume_response_body_chunk, produce_response_body_chunk = create_channel(1)
         request_complete = False
-        response_started = False
+        response_started = create_event()
         response_complete = create_event()
+        app_crashed = create_event()
+        app_exception: typing.Optional[Exception] = None
 
         headers = [] if headers is None else headers
         stream = ByteStream(b"") if stream is None else stream
 
         request_body_chunks = stream.__aiter__()
 
         async def receive() -> dict:
-            nonlocal request_complete, response_complete
+            nonlocal request_complete
 
             if request_complete:
                 await response_complete.wait()
@@ -120,38 +210,76 @@ async def receive() -> dict:
             return {"type": "http.request", "body": body, "more_body": True}
 
         async def send(message: dict) -> None:
-            nonlocal status_code, response_headers, body_parts
-            nonlocal response_started, response_complete
+            nonlocal status_code, response_headers
 
             if message["type"] == "http.response.start":
-                assert not response_started
+                assert not response_started.is_set()
 
                 status_code = message["status"]
                 response_headers = message.get("headers", [])
-                response_started = True
+                response_started.set()
 
             elif message["type"] == "http.response.body":
                 assert not response_complete.is_set()
                 body = message.get("body", b"")
                 more_body = message.get("more_body", False)
 
                 if body and method != b"HEAD":
-                    body_parts.append(body)
+                    await produce_response_body_chunk(body)
 
                 if not more_body:
                     response_complete.set()
 
-        try:
-            await self.app(scope, receive, send)
-        except Exception:
-            if self.raise_app_exceptions or not response_complete:
-                raise
+        async def run_app() -> None:
+            nonlocal app_exception
+            try:
+                await self.app(scope, receive, send)
+            except Exception as exc:
+                app_exception = exc
+                app_crashed.set()
+
+        background = create_background_task(run_app)
+        await background.start()
+
+        await run_until_first_complete(app_crashed.wait, response_started.wait)
 
-        assert response_complete.is_set()
+        if app_crashed.is_set():
+            assert app_exception is not None
+            await background.aclose(app_exception)
+            if self.raise_app_exceptions or not response_started.is_set():
+                raise app_exception
+
+        assert response_started.is_set()
         assert status_code is not None
         assert response_headers is not None
 
-        stream = ByteStream(b"".join(body_parts))
+        async def aiter_response_body_chunks() -> typing.AsyncIterator[bytes]:
+            chunk = b""
+
+            async def consume_chunk() -> None:
+                nonlocal chunk
+                chunk = await consume_response_body_chunk()
+
+            while True:
+                await run_until_first_complete(
+                    app_crashed.wait, consume_chunk, response_complete.wait
+                )
+
+                if app_crashed.is_set():
+                    assert app_exception is not None
+                    if self.raise_app_exceptions:
+                        raise app_exception
+                    break
+
+                yield chunk
+
+                if response_complete.is_set():
+                    break
+
+        async def aclose() -> None:
+            await background.aclose(app_exception)
+
+        stream = AsyncIteratorStream(aiter_response_body_chunks(), close_func=aclose)
 
         return (b"HTTP/1.1", status_code, b"", response_headers, stream)
 

tests/test_asgi.py

@@ -1,6 +1,7 @@
 import pytest
 
 import httpx
+from .concurrency import sleep
 
 
 async def hello_world(scope, receive, send):
@@ -35,7 +36,8 @@ async def raise_exc_after_response(scope, receive, send):
     headers = [(b"content-type", "text/plain"), (b"content-length", str(len(output)))]
 
     await send({"type": "http.response.start", "status": status, "headers": headers})
-    await send({"type": "http.response.body", "body": output})
+    await send({"type": "http.response.body", "body": output, "more_body": True})
+    await sleep(0.001)  # Let the transport detect that the response has started.
     raise ValueError()
 
 
@@ -99,3 +101,29 @@ async def read_body(scope, receive, send):
     response = await client.post("http://www.example.org/", data=b"example")
     assert response.status_code == 200
     assert disconnect
+
+
+@pytest.mark.asyncio
+async def test_asgi_streaming():
+    client = httpx.AsyncClient(app=hello_world)
+    async with client.stream("GET", "http://www.example.org/") as response:
+        assert response.status_code == 200
+        text = "".join([chunk async for chunk in response.aiter_text()])
+        assert text == "Hello, World!"
+
+
+@pytest.mark.asyncio
+async def test_asgi_streaming_exc():
+    client = httpx.AsyncClient(app=raise_exc)
+    with pytest.raises(ValueError):
+        async with client.stream("GET", "http://www.example.org/"):
+            pass  # pragma: no cover
+
+
+@pytest.mark.asyncio
+async def test_asgi_streaming_exc_after_response():
+    client = httpx.AsyncClient(app=raise_exc_after_response)
+    async with client.stream("GET", "http://www.example.org/") as response:
+        with pytest.raises(ValueError):
+            async for _ in response.aiter_bytes():
+                pass  # pragma: no cover