base.py

import typing

import anyio

from starlette.requests import Request
from starlette.responses import Response, StreamingResponse
from starlette.types import ASGIApp, Receive, Scope, Send

RequestResponseEndpoint = typing.Callable[[Request], typing.Awaitable[Response]]
DispatchFunction = typing.Callable[
    [Request, RequestResponseEndpoint], typing.Awaitable[Response]
]


class BaseHTTPMiddleware:
    def __init__(self, app: ASGIApp, dispatch: DispatchFunction = None) -> None:
        self.app = app
        self.dispatch_func = self.dispatch if dispatch is None else dispatch

    async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
        if scope["type"] != "http":
            await self.app(scope, receive, send)
            return

        async def call_next(request: Request) -> Response:
            send_stream, recv_stream = anyio.create_memory_object_stream()

            async def coro() -> None:
                async with send_stream:
                    await self.app(scope, request.receive, send_stream.send)

            task_group.start_soon(coro)

            try:
                message = await recv_stream.receive()
            except anyio.EndOfStream:
                raise RuntimeError("No response returned.") from None

            assert message["type"] == "http.response.start"

            async def body_stream() -> typing.AsyncGenerator[bytes, None]:
                async with recv_stream:
                    async for message in recv_stream:
                        assert message["type"] == "http.response.body"
                        yield message.get("body", b"")

            response = StreamingResponse(
                status_code=message["status"], content=body_stream()
            )
            response.raw_headers = message["headers"]
            return response

        try:
            async with anyio.create_task_group() as task_group:
                request = Request(scope, receive=receive)
                response = await self.dispatch_func(request, call_next)
                await response(scope, receive, send)
                task_group.cancel_scope.cancel()
        except anyio.ExceptionGroup as exc_group:
            first, *rest = exc_group.exceptions
            if str(first) == "No response returned.":
                # "No response returned", and we got an exception group, meaning
                # the app raised an exception.
                # Raise the latter in priority, so that we:
                # a/ match the behavior before 0.16.0 (see regression #1255).
                # b/ keep `anyio` as an implementation detail internally.
                assert len(rest) == 1
                (exc,) = rest
                raise exc from None
            else:
                # User may have used `anyio` directly and triggered an
                # `ExceptionGroup` which they may want to manage themselves.
                # Let it be.
                raise

    async def dispatch(
        self, request: Request, call_next: RequestResponseEndpoint
    ) -> Response:
        raise NotImplementedError()  # pragma: no cover