Make ASGITransport streamable

[jonathanslenders]

Right now, httpx's ASGITransport does not operate on a stream, it buffers the whole response internally before proceeding.

I noticed there is this related issue on the httpx repo: encode/httpx#3391

But independent from that, I successfully implemented a StreamingASGITransport which does not buffer internally.
The API is slightly different: the creation of the transport requires the use of a context manager because a task group is needed for running the ASGI endpoint. The code is below.

Feel free to integrate into httpx2 or to ignore. Leaving it here in case it's useful to somebody. I can create a PR if you'd like (but I won't have time to write docs if that's needed.)

from __future__ import annotations

import typing
from collections.abc import AsyncGenerator, Awaitable, Callable, MutableMapping
from contextlib import asynccontextmanager

from anyio import (
    BrokenResourceError,
    Event,
    create_memory_object_stream,
    create_task_group,
)
from anyio.abc import TaskGroup
from anyio.streams.memory import MemoryObjectReceiveStream
from httpx._models import Request, Response
from httpx._transports.base import AsyncBaseTransport
from httpx._types import AsyncByteStream
from typing_extensions import Self

__all__ = ["StreamingASGITransport"]

_Message = MutableMapping[str, typing.Any]
_Receive = Callable[[], Awaitable[_Message]]
_Send = Callable[[MutableMapping[str, typing.Any]], Awaitable[None]]
_ASGIApp = Callable[[MutableMapping[str, typing.Any], _Receive, _Send], Awaitable[None]]



class ASGIResponseStream(AsyncByteStream):
    def __init__(self, response_receiver: MemoryObjectReceiveStream[bytes]) -> None:
        self._receiver = response_receiver

    async def __aiter__(self) -> AsyncGenerator[bytes]:
        with self._receiver:
            async for data in self._receiver:
                yield data


class StreamingASGITransport(AsyncBaseTransport):
    """
    A custom AsyncTransport that handles sending requests directly to an ASGI app.

    ```python
    async with StreamingSGITransport.create(
        app=app,
        root_path="/submount",
        client=("1.2.3.4", 123)
    ) as transport:
        client = httpx.AsyncClient(transport=transport)
    ```

    Arguments:

    * `app` - The ASGI application.
    * `root_path` - The root path on which the ASGI application should be mounted.
    * `client` - A two-tuple indicating the client IP and port of incoming requests.
    ```
    """

    def __init__(
        self,
        app: _ASGIApp,
        root_path: str,
        client: tuple[str, int],
        tg: TaskGroup,
    ) -> None:
        self.app = app
        self.root_path = root_path
        self.client = client
        self._tg = tg

    @classmethod
    @asynccontextmanager
    async def create(
        cls,
        app: _ASGIApp,
        root_path: str = "",
        client: tuple[str, int] = ("127.0.0.1", 123),
    ) -> AsyncGenerator[Self]:
        async with create_task_group() as tg:
            yield cls(
                app=app,
                root_path=root_path,
                client=client,
                tg=tg,
            )
            tg.cancel_scope.cancel()

    async def handle_async_request(
        self,
        request: Request,
    ) -> Response:
        assert isinstance(request.stream, AsyncByteStream)

        # ASGI scope.
        scope = {
            "type": "http",
            "asgi": {"version": "3.0"},
            "http_version": "1.1",
            "method": request.method,
            "headers": [(k.lower(), v) for (k, v) in request.headers.raw],
            "scheme": request.url.scheme,
            "path": request.url.path,
            "raw_path": request.url.raw_path.split(b"?")[0],
            "query_string": request.url.query,
            "server": (request.url.host, request.url.port),
            "client": self.client,
            "root_path": self.root_path,
        }

        # Request.
        request_body_chunks = request.stream.__aiter__()
        request_complete = False

        # Response.
        status_code = None
        response_headers = None
        response_started = Event()
        response_complete = Event()

        response_sender, response_receiver = create_memory_object_stream[bytes]()

        # ASGI callables.

        async def receive() -> dict[str, typing.Any]:
            nonlocal request_complete

            if request_complete:
                await response_complete.wait()
                return {"type": "http.disconnect"}

            try:
                body = await request_body_chunks.__anext__()
            except StopAsyncIteration:
                request_complete = True
                return {"type": "http.request", "body": b"", "more_body": False}
            return {"type": "http.request", "body": body, "more_body": True}

        async def send(message: MutableMapping[str, typing.Any]) -> None:
            nonlocal status_code, response_headers, response_started

            if message["type"] == "http.response.start":
                assert not response_started.is_set()

                status_code = message["status"]
                response_headers = message.get("headers", [])
                response_started.set()

            elif message["type"] == "http.response.body":
                assert not response_complete.is_set()

                try:
                    body = message.get("body", b"")
                    more_body = message.get("more_body", False)

                    if body and request.method != "HEAD":
                        await response_sender.send(body)

                    if not more_body:
                        response_sender.close()
                        response_complete.set()
                except BrokenResourceError:
                    # Raise OSError when the client disconnects:
                    # https://asgi.readthedocs.io/en/stable/specs/www.html#disconnected-client-send-exception
                    raise OSError("Client disconnected")

        async def run_app() -> None:
            with response_sender:
                nonlocal status_code, response_headers
                try:
                    await self.app(scope, receive, send)
                except Exception:  # noqa: PIE-786
                    # TODO: Log exception!

                    response_started.set()
                    response_complete.set()
                    if status_code is None:
                        status_code = 500
                    if response_headers is None:
                        response_headers = {}

                assert response_complete.is_set()
                assert status_code is not None
                assert response_headers is not None

        self._tg.start_soon(run_app)

        stream = ASGIResponseStream(response_receiver)

        # Wait for the headers to be set.
        await response_started.wait()
        assert status_code is not None

        return Response(status_code, headers=response_headers, stream=stream)