fix: @inject works on HTTPEndpoint / WebSocketEndpoint methods

README.md

@@ -76,6 +76,40 @@ container.validate()  # optional fail-fast; must come after setup_di registers i
 
 Call `setup_di` once, after creating the app and before it starts serving — it installs middleware, and Starlette does not allow middleware to be added after startup.
 
+`@inject` works the same on the methods of a class-based endpoint. Decorate the handler method, not the class; `self` and any arguments Starlette passes after the connection are forwarded unchanged, so `WebSocketEndpoint.on_receive` and `on_disconnect` inject too:
+
+```python
+from starlette.endpoints import HTTPEndpoint, WebSocketEndpoint
+from starlette.routing import Route, WebSocketRoute
+from starlette.websockets import WebSocket
+
+
+class Users(HTTPEndpoint):
+    @inject
+    async def get(
+        self,
+        request: Request,
+        service: typing.Annotated[UserService, FromDI(Dependencies.user_service)],
+    ) -> JSONResponse:
+        return JSONResponse({"debug": service.settings.debug})
+
+
+class Echo(WebSocketEndpoint):
+    encoding = "text"
+
+    @inject
+    async def on_receive(
+        self,
+        websocket: WebSocket,
+        data: str,
+        service: typing.Annotated[UserService, FromDI(Dependencies.user_service)],
+    ) -> None:
+        await websocket.send_text(data)
+
+
+app = Starlette(routes=[Route("/users", Users), WebSocketRoute("/echo", Echo)])
+```
+
 An HTTP request opens a `Scope.REQUEST` child container; a WebSocket connection opens a `Scope.SESSION` one, both built by the middleware before your handler runs. The connection `starlette.requests.Request` / `starlette.websockets.WebSocket` are resolvable within DI via the pre-built `starlette_request_provider` / `starlette_websocket_provider` context providers. The instance a provider receives is backed by the same ASGI scope as your handler's connection but is a distinct object: read `method` / `url` / `headers` / `state` from it, not the request body.
 
 ## API
@@ -84,7 +118,7 @@ An HTTP request opens a `Scope.REQUEST` child container; a WebSocket connection
 |---|---|
 | `setup_di(app, container)` | Registers the container on `app.state`, composes the lifespan (opens/closes the container), and installs the middleware that builds a per-connection child container; returns the container |
 | `FromDI(dependency)` | Inert marker (used with `@inject`) that resolves a provider or type from the per-connection child container |
-| `inject(handler)` | Decorator for an `async def` handler taking a `Request` or `WebSocket`; resolves its `FromDI`-annotated parameters |
+| `inject(handler)` | Decorator for an `async def` handler taking a `Request` or `WebSocket`, either a function endpoint or a method of an `HTTPEndpoint` / `WebSocketEndpoint` subclass; resolves its `FromDI`-annotated parameters |
 | `fetch_di_container(app)` | Returns the root `Container` stored on `app.state` |
 | `starlette_request_provider` | `ContextProvider` for `starlette.requests.Request` (`REQUEST` scope), auto-registered |
 | `starlette_websocket_provider` | `ContextProvider` for `starlette.websockets.WebSocket` (`SESSION` scope), auto-registered |

docs/adr/0001-pure-asgi-middleware.md

@@ -1,27 +1,10 @@
 # The container is opened in pure ASGI middleware, not `BaseHTTPMiddleware`
 
-**Decision:** `_DIMiddleware` is a plain ASGI callable
-(`async def __call__(self, scope, receive, send)`). We will not implement it as a Starlette
-`BaseHTTPMiddleware` subclass.
-
-`BaseHTTPMiddleware` is the ergonomic choice and the one Starlette's own tutorial reaches for
-first: it hands you a `Request` and a `call_next`, and the child container would open and close
-around a single `await call_next(request)`. It is rejected because it runs the downstream app in a
-separate anyio task. `contextvars` set on one side of that boundary are not visible on the other,
-and a DI container whose scope does not survive into the endpoint is not a DI container. Starlette's
-own middleware documentation carries the warning; this is not a subtlety we discovered.
-
-Two further consequences follow from the plain-callable form, and both are load-bearing rather than
-incidental. WebSockets are reachable at all: `BaseHTTPMiddleware` handles `http` only, so a
-`Scope.SESSION` child container for a WebSocket connection would need a second, differently shaped
-mechanism. And the ASGI `scope` dict is in hand, which is how the child container reaches `@inject`
-without a `contextvar` in the first place.
-
-The cost is that the middleware is written against the raw ASGI three-argument protocol and must
-construct its own `Request` / `WebSocket`, pass non-connection scope types (`lifespan`) straight
-through, and close the child container on the exception path itself. That is roughly fifteen lines,
-paid once.
-
-**Revisit trigger:** Starlette makes `BaseHTTPMiddleware` share a context with the downstream app —
-or ships a supported middleware base that does — *and* it covers `websocket` scopes. Both halves
-have to land: either one alone leaves this integration writing the raw protocol anyway.
+**Decision:** `_DIMiddleware` is a plain ASGI callable, not a `BaseHTTPMiddleware` subclass.
+`BaseHTTPMiddleware` is the ergonomic choice, but it runs the downstream app in a separate anyio
+task, so `contextvars` set around `call_next` are not visible in the endpoint, and it handles `http`
+only, so a `Scope.SESSION` child for a WebSocket would need a second mechanism. The plain callable
+costs about fifteen lines of raw-protocol handling and in return has the ASGI `scope` dict in hand,
+which is how the child container reaches `@inject` without a `contextvar`. **Revisit trigger:**
+Starlette ships a supported middleware base that shares context with the downstream app *and*
+covers `websocket` scopes; either half alone leaves this integration writing the raw protocol.

docs/adr/0002-child-container-stays-internal.md

@@ -1,30 +1,13 @@
 # The per-connection child container has no public accessor
 
-**Decision:** the child container is reachable only through `@inject` + `FromDI`. The ASGI scope key
-it lives under stays private, and we will not add a `fetch_di_child_container(connection)` (or
-equivalent) to the public surface. `fetch_di_container(app)` returns the root container and is
-deliberately the only container accessor a user gets.
-
-The obvious counterpart to `fetch_di_container(app)` is a per-connection twin, and it would be two
-lines. It was declined at v1 as unneeded, and a later change turned that from a preference into a
-constraint. The middleware now deletes its scope entry in a `finally` when the connection ends,
-because the container's context holds the connection, the connection owns the ASGI `scope` dict, and
-the `scope` dict held the container — a cycle per request, leaving a finished request reclaimable
-only by the garbage collector rather than by refcounting. Clearing the entry took that from 34
-cyclic objects per request to zero.
-
-A public accessor advertises the entry as a thing callers may hold. The moment one is handed out and
-kept — stored on an object, closed over by a background task, read after the response — the entry is
-either still present and the cycle is back, or it is gone and the accessor raises. There is no
-version of the accessor that is both safe and useful, because the lifetime it exposes is strictly
-shorter than the object a caller would want to attach it to. The bound is the point: nothing may
-read the entry after the middleware's `async with` block exits.
-
-The one real gap this leaves is class-based `HTTPEndpoint` / `WebSocketEndpoint`, which `@inject`
-does not cover. That is a missing decorator path, not a missing accessor, and is tracked as its own
-work.
-
-**Revisit trigger:** a use for the child container appears that `@inject` genuinely cannot serve —
-not a class-based endpoint, which wants its own injection path, but something outside the
-connection's own call stack. At that point the lifetime question above has to be answered first, and
-the answer is what the accessor's contract would be.
+**Decision:** the child container is reachable only through `@inject` + `FromDI`; the ASGI scope
+key it lives under stays private and there is no `fetch_di_child_container(connection)`. The
+middleware deletes its scope entry in a `finally` when the connection ends, because the container's
+context holds the connection, the connection owns the `scope` dict, and the dict held the container:
+a cycle per request that left finished requests to the garbage collector. An accessor advertises the
+entry as something a caller may hold, and any holder that outlives the connection either revives the
+cycle or reads a deleted entry, so there is no version of it that is both safe and useful. Class-based
+`HTTPEndpoint` / `WebSocketEndpoint` were once the gap this left; `@inject` now binds as a method
+(modern-python/modern-di-starlette#28), so it was a missing decorator path, not a missing accessor.
+**Revisit trigger:** a use for the child container appears outside the connection's own call stack;
+the lifetime question above has to be answered first, and the answer is the accessor's contract.

modern_di_starlette/main.py

@@ -86,11 +86,21 @@ def setup_di(app: Starlette, container: Container) -> Container:
 FromDI = integrations.from_di
 
 
+def _find_connection(handler_name: str, args: tuple[typing.Any, ...]) -> Request | WebSocket:
+    for arg in args:
+        if isinstance(arg, Request | WebSocket):
+            return arg
+    msg = f"@inject requires {handler_name} to receive a Request or WebSocket as a positional argument."
+    raise TypeError(msg)
+
+
 def inject(func: typing.Callable[..., typing.Awaitable[T]]) -> typing.Callable[..., typing.Awaitable[T]]:
     markers = integrations.parse_markers(func)
+    handler_name = getattr(func, "__qualname__", repr(func))
 
     @functools.wraps(func)
-    async def wrapper(connection: Request | WebSocket) -> T:
+    async def wrapper(*args: typing.Any, **kwargs: typing.Any) -> T:  # noqa: ANN401
+        connection = _find_connection(handler_name, args)
         try:
             child_container: Container = connection.scope[_CONTAINER_SCOPE_KEY]
         except KeyError:
@@ -100,6 +110,6 @@ async def wrapper(connection: Request | WebSocket) -> T:
                 "before using @inject."
             )
             raise RuntimeError(msg) from None
-        return await func(connection, **integrations.resolve_markers(child_container, markers))
+        return await func(*args, **kwargs, **integrations.resolve_markers(child_container, markers))
 
     return wrapper

tests/test_endpoints.py

@@ -0,0 +1,95 @@
+import typing
+
+import pytest
+from starlette import status
+from starlette.applications import Starlette
+from starlette.endpoints import HTTPEndpoint, WebSocketEndpoint
+from starlette.requests import Request
+from starlette.responses import PlainTextResponse
+from starlette.routing import Route, WebSocketRoute
+from starlette.testclient import TestClient
+from starlette.websockets import WebSocket
+
+from modern_di_starlette import FromDI, inject
+from tests.dependencies import Dependencies, DependentCreator, SimpleCreator
+
+
+def test_http_endpoint_method_resolves_markers(client: TestClient, app: Starlette) -> None:
+    class Endpoint(HTTPEndpoint):
+        @inject
+        async def get(
+            self,
+            request: Request,
+            app_factory_instance: typing.Annotated[SimpleCreator, FromDI(SimpleCreator)],
+            request_factory_instance: typing.Annotated[DependentCreator, FromDI(Dependencies.request_factory)],
+            method: typing.Annotated[str, FromDI(Dependencies.request_method)],
+        ) -> PlainTextResponse:
+            assert isinstance(self, Endpoint)
+            assert isinstance(request, Request)
+            assert isinstance(app_factory_instance, SimpleCreator)
+            assert isinstance(request_factory_instance, DependentCreator)
+            assert request_factory_instance.dep1 is not app_factory_instance
+            return PlainTextResponse(method)
+
+    app.router.routes.append(Route("/", Endpoint))
+    response = client.get("/")
+    assert response.status_code == status.HTTP_200_OK
+    assert response.text == "GET"
+
+
+def test_websocket_endpoint_hooks_resolve_markers(client: TestClient, app: Starlette) -> None:
+    seen: list[str] = []
+
+    class Endpoint(WebSocketEndpoint):
+        encoding = "text"
+
+        @inject
+        async def on_connect(
+            self,
+            websocket: WebSocket,
+            session_factory_instance: typing.Annotated[DependentCreator, FromDI(Dependencies.session_factory)],
+            path: typing.Annotated[str, FromDI(Dependencies.websocket_path)],
+        ) -> None:
+            assert isinstance(self, Endpoint)
+            assert isinstance(session_factory_instance, DependentCreator)
+            await websocket.accept()
+            await websocket.send_text(path)
+
+        @inject
+        async def on_receive(
+            self,
+            websocket: WebSocket,
+            data: str,
+            app_factory_instance: typing.Annotated[SimpleCreator, FromDI(SimpleCreator)],
+        ) -> None:
+            assert isinstance(websocket, WebSocket)
+            await websocket.send_text(f"{data}:{app_factory_instance.dep1}")
+
+        @inject
+        async def on_disconnect(
+            self,
+            websocket: WebSocket,
+            close_code: int,
+            path: typing.Annotated[str, FromDI(Dependencies.websocket_path)],
+        ) -> None:
+            assert isinstance(websocket, WebSocket)
+            seen.append(f"{path}:{close_code}")
+
+    app.router.routes.append(WebSocketRoute("/ws", Endpoint))
+    with client.websocket_connect("/ws") as websocket:
+        assert websocket.receive_text() == "/ws"
+        websocket.send_text("ping")
+        assert websocket.receive_text() == "ping:original"
+    assert seen == [f"/ws:{status.WS_1000_NORMAL_CLOSURE}"]
+
+
+async def test_inject_without_connection_argument_raises_clear_error() -> None:
+    @inject
+    async def handler(
+        value: str,
+        app_factory_instance: typing.Annotated[SimpleCreator, FromDI(SimpleCreator)],
+    ) -> str:
+        return f"{value}:{app_factory_instance.dep1}"  # pragma: no cover -- TypeError precedes this call
+
+    with pytest.raises(TypeError, match="Request or WebSocket"):
+        await handler("value")

tests/test_middleware.py

@@ -1,11 +1,14 @@
 import gc
 import sys
 
+import pytest
 from modern_di import Container, Scope
 from starlette import status
 from starlette.applications import Starlette
+from starlette.endpoints import HTTPEndpoint
 from starlette.requests import Request
 from starlette.responses import PlainTextResponse
+from starlette.routing import Route
 from starlette.testclient import TestClient
 from starlette.types import Scope as ASGIScope
 
@@ -26,7 +29,8 @@ def endpoint(request: Request) -> PlainTextResponse:
     assert client.get("/").status_code == status.HTTP_200_OK
 
 
-def test_finished_request_leaves_no_cyclic_garbage(client: TestClient, app: Starlette) -> None:
+@pytest.mark.parametrize("class_based", [False, True], ids=["function", "HTTPEndpoint"])
+def test_finished_request_leaves_no_cyclic_garbage(client: TestClient, app: Starlette, class_based: bool) -> None:
     """INVARIANT: a completed connection leaves no reference cycle behind.
 
     Broken by anything that lets the child container stay reachable from the ASGI scope past the
@@ -52,7 +56,11 @@ def endpoint(request: Request) -> PlainTextResponse:
         last_scope[:] = [request.scope]
         return PlainTextResponse("ok")
 
-    app.add_route("/", endpoint)
+    class Endpoint(HTTPEndpoint):
+        async def get(self, request: Request) -> PlainTextResponse:
+            return endpoint(request)
+
+    app.router.routes.append(Route("/", Endpoint if class_based else endpoint))
     for _ in range(5):  # let one-time allocations settle before measuring
         client.get("/")