restore original ASGITransport class code and split streaming-enabled version into separate ASGIStreamingTransport class

Author: jhominal

httpx/_transports/__init__.py

@@ -6,6 +6,7 @@
 
 __all__ = [
     "ASGITransport",
+    "ASGIStreamingTransport",
     "AsyncBaseTransport",
     "BaseTransport",
     "AsyncHTTPTransport",

httpx/_transports/asgi.py

@@ -33,7 +33,7 @@
     [typing.MutableMapping[str, typing.Any], _Receive, _Send], typing.Awaitable[None]
 ]
 
-__all__ = ["ASGITransport"]
+__all__ = ["ASGITransport", "ASGIStreamingTransport"]
 
 
 def is_running_trio() -> bool:
@@ -98,6 +98,141 @@ def get_end_of_stream_error_type() -> type[anyio.EndOfStream | trio.EndOfChannel
 
 
 class ASGIResponseStream(AsyncByteStream):
+    def __init__(self, body: list[bytes]) -> None:
+        self._body = body
+
+    async def __aiter__(self) -> typing.AsyncIterator[bytes]:
+        yield b"".join(self._body)
+
+
+class ASGITransport(AsyncBaseTransport):
+    """
+    A custom AsyncTransport that handles sending requests directly to an ASGI app.
+
+    ```python
+    transport = httpx.ASGITransport(
+        app=app,
+        root_path="/submount",
+        client=("1.2.3.4", 123)
+    )
+    client = httpx.AsyncClient(transport=transport)
+    ```
+
+    Arguments:
+
+    * `app` - The ASGI application.
+    * `raise_app_exceptions` - Boolean indicating if exceptions in the application
+       should be raised. Default to `True`. Can be set to `False` for use cases
+       such as testing the content of a client 500 response.
+    * `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,
+        raise_app_exceptions: bool = True,
+        root_path: str = "",
+        client: tuple[str, int] = ("127.0.0.1", 123),
+    ) -> None:
+        self.app = app
+        self.raise_app_exceptions = raise_app_exceptions
+        self.root_path = root_path
+        self.client = client
+
+    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
+        body_parts = []
+        response_started = False
+        response_complete = create_event()
+
+        # 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: typing.MutableMapping[str, typing.Any]) -> None:
+            nonlocal status_code, response_headers, response_started
+
+            if message["type"] == "http.response.start":
+                assert not response_started
+
+                status_code = message["status"]
+                response_headers = message.get("headers", [])
+                response_started = True
+
+            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 request.method != "HEAD":
+                    body_parts.append(body)
+
+                if not more_body:
+                    response_complete.set()
+
+        try:
+            await self.app(scope, receive, send)
+        except Exception:  # noqa: PIE-786
+            if self.raise_app_exceptions:
+                raise
+
+            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
+
+        stream = ASGIResponseStream(body_parts)
+
+        return Response(status_code, headers=response_headers, stream=stream)
+
+
+class ASGIStreamingResponseStream(AsyncByteStream):
     def __init__(
         self,
         ignore_body: bool,
@@ -124,18 +259,19 @@ async def aclose(self) -> None:
         await self._asgi_generator.aclose()
 
 
-class ASGITransport(AsyncBaseTransport):
+class ASGIStreamingTransport(AsyncBaseTransport):
     """
-    A custom AsyncTransport that handles sending requests directly to an ASGI app.
+    An equivalent of ASGITransport that operates by running app in a sub-task and
+    streaming response events as soon as they arrive.
 
-    ```python
-    transport = httpx.ASGITransport(
-        app=app,
-        root_path="/submount",
-        client=("1.2.3.4", 123)
-    )
-    client = httpx.AsyncClient(transport=transport)
-    ```
+    It is used in the same way, with the same arguments having the same signification,
+    as ASGITransport.
+
+    The main observable differences between the two implementations will be as follows:
+     * As the application callable is invoked in a sub-task, any context variables that
+       are set by the app will not propagate to the caller;
+     * The streaming mode of operation means that a response will generally be returned
+       to the AsyncClient caller before the application has fully run;
 
     Arguments:
 
@@ -145,9 +281,6 @@ class ASGITransport(AsyncBaseTransport):
        such as testing the content of a client 500 response.
     * `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.
-    * `streaming` - Set to `True` to enable streaming of response content. Default to
-      `False`, as activating this feature means that the ASGI `app` will run in a
-      sub-task, which has observable side effects for context variables.
     ```
     """
 
@@ -157,14 +290,11 @@ def __init__(
         raise_app_exceptions: bool = True,
         root_path: str = "",
         client: tuple[str, int] = ("127.0.0.1", 123),
-        *,
-        streaming: bool = False,
     ) -> None:
         self.app = app
         self.raise_app_exceptions = raise_app_exceptions
         self.root_path = root_path
         self.client = client
-        self.streaming = streaming
 
     async def handle_async_request(
         self,
@@ -177,7 +307,7 @@ async def handle_async_request(
                 return Response(
                     status_code=message["status"],
                     headers=message.get("headers", []),
-                    stream=ASGIResponseStream(
+                    stream=ASGIStreamingResponseStream(
                         ignore_body=request.method == "HEAD",
                         asgi_generator=asgi_generator,
                     ),
@@ -214,9 +344,8 @@ async def _stream_asgi_messages(
         response_complete = create_event()
 
         # ASGI response messages stream
-        stream_size = 0 if self.streaming else float("inf")
         response_message_send_stream, response_message_recv_stream = (
-            create_memory_object_stream(stream_size)
+            create_memory_object_stream(0)
         )
 
         # ASGI app exception
@@ -256,11 +385,8 @@ async def run_app() -> None:
 
         async with contextlib.AsyncExitStack() as exit_stack:
             exit_stack.callback(response_complete.set)
-            if self.streaming:
-                task_group = await exit_stack.enter_async_context(create_task_group())
-                task_group.start_soon(run_app)
-            else:
-                await run_app()
+            task_group = await exit_stack.enter_async_context(create_task_group())
+            task_group.start_soon(run_app)
 
             async with response_message_recv_stream:
                 try: