Initial implementation of a Hook system to build middlewares.

main.py

from __future__ import annotations

from typing import TYPE_CHECKING, Awaitable, Callable, TypedDict, TypeVar

from starlette.datastructures import MutableHeaders
from typing_extensions import NotRequired

if TYPE_CHECKING:
    from asgiref.typing import (
        ASGI3Application,
        ASGIReceiveCallable,
        ASGISendCallable,
        ASGISendEvent,
        HTTPResponseBodyEvent,
        HTTPResponseStartEvent,
        Scope,
    )


_T = TypeVar("_T")
HookCallable = Callable[["Scope", _T], Awaitable[None]]


class Hook(TypedDict):
    http_response_start: NotRequired[HookCallable[HTTPResponseStartEvent]]
    http_response_body: NotRequired[HookCallable[HTTPResponseBodyEvent]]


class HookMiddleware:
    def __init__(self, app: ASGI3Application, hook: Hook) -> None:
        self.app = app
        self.hook = hook

    async def __call__(
        self, scope: Scope, receive: ASGIReceiveCallable, send: ASGISendCallable
    ) -> None:
        if scope["type"] == "http":

            async def send_wrapper(event: ASGISendEvent) -> None:
                if event["type"] == "http.response.start":
                    if "http_response_start" in self.hook:
                        await self.hook["http_response_start"](scope, event)
                elif event["type"] == "http.response.body":
                    if "http_response_body" in self.hook:
                        await self.hook["http_response_body"](scope, event)
                await send(event)

            return await self.app(scope, receive, send_wrapper)
        elif scope["type"] == "websocket":
            return await self.app(scope, receive, send)
        else:
            return await self.app(scope, receive, send)


async def http_response_start(scope: Scope, message: HTTPResponseStartEvent) -> None:
    headers = MutableHeaders(scope=message)
    for key, value in [("content-type", "text/html"), ("accept", "text/html")]:
        headers.append(key, value)


async def app(scope, receive, send) -> None:
    await send({"type": "http.response.start", "status": 200, "headers": []})
    await send({"type": "http.response.body", "body": b"Hello, world!"})


application = HookMiddleware(app=app, hook={"http_response_start": http_response_start})