fix: make the request container usable from every request hook

CONTEXT.md

@@ -27,8 +27,9 @@ _Avoid_: app container — reads as Flask's app object, and its lifetime is the
 Flask app context.
 
 **Child container**:
-The `Scope.REQUEST` container `before_request` builds per request and stashes on `flask.g`;
-`inject` resolves from it and `teardown_appcontext` closes it. Its lifetime is bounded by Flask's
+The `Scope.REQUEST` container built per request, on first use or in `before_request`, whichever
+comes first, and stashed on `flask.g`; `inject` resolves from it and `teardown_appcontext` closes
+it after every other `teardown_appcontext` handler. Its lifetime is bounded by Flask's
 **app context**, not by Flask's request context: an app context pushed without a request (a CLI
 command, a bare `app.app_context()`) fires the same teardown with no child ever built, which is
 why the close is guarded.

README.md

@@ -31,7 +31,7 @@ uv add modern-di-flask      # or: pip install modern-di-flask
 
 ## Usage
 
-Flask has no dependency-injection system of its own, so `modern-di-flask` pairs an `@inject` decorator with inert `FromDI` markers (there is no `Depends`). `setup_di` installs a `before_request`/`teardown_appcontext` pair that builds a per-request `Scope.REQUEST` child container and closes it once the request finishes. Resolution is sync, also for `async def` views, and the child container is closed with `close_sync()`.
+Flask has no dependency-injection system of its own, so `modern-di-flask` pairs an `@inject` decorator with inert `FromDI` markers (there is no `Depends`). `setup_di` installs a `before_request`/`teardown_appcontext` pair that builds a per-request `Scope.REQUEST` child container and closes it once the request finishes. The child is also built on first use, so request hooks registered before `setup_di` can use it, and it is closed after every `teardown_appcontext` handler. Resolution is sync, also for `async def` views, and the child container is closed with `close_sync()`.
 
 ```python
 import typing

modern_di_flask/main.py

@@ -14,6 +14,7 @@
 # constants keep each writer and reader in provable agreement.
 _ROOT_CONTAINER_KEY = "modern_di_container"
 _CHILD_CONTAINER_ATTR = "modern_di_request_container"
+_CHILD_REQUEST_ATTR = "modern_di_request"
 _NO_CHILD_CONTAINER_MSG = (
     "No modern-di container found for this request. Call setup_di(app, container) "
     "so requests pass through the modern-di request hooks before using @inject."
@@ -25,12 +26,33 @@ def fetch_di_container(app: Flask) -> Container:
 
 
 def _fetch_child_container() -> Container:
+    if flask.has_request_context():
+        return _request_child_container()
     try:
         return typing.cast(Container, getattr(g, _CHILD_CONTAINER_ATTR))
     except AttributeError:
         raise RuntimeError(_NO_CHILD_CONTAINER_MSG) from None
 
 
+def _request_child_container() -> Container:
+    connection = flask.request._get_current_object()  # noqa: SLF001  # ty: ignore[unresolved-attribute]
+    if getattr(g, _CHILD_REQUEST_ATTR, None) is connection:
+        return typing.cast(Container, getattr(g, _CHILD_CONTAINER_ATTR))
+    try:
+        container = fetch_di_container(flask.current_app)
+    except KeyError:
+        raise RuntimeError(_NO_CHILD_CONTAINER_MSG) from None
+    match = integrations.bind(flask_request_provider, connection)
+    child = container.build_child_container(scope=match.scope, context=match.context)
+    setattr(g, _CHILD_CONTAINER_ATTR, child)
+    setattr(g, _CHILD_REQUEST_ATTR, connection)
+    return child
+
+
+def _enter_request() -> None:
+    _request_child_container()
+
+
 def _close_request(_exception: BaseException | None) -> None:
     child: Container | None = getattr(g, _CHILD_CONTAINER_ATTR, None)
     if child is not None:
@@ -41,14 +63,11 @@ def setup_di(app: Flask, container: Container, *, auto_inject: bool = False) ->
     app.extensions[_ROOT_CONTAINER_KEY] = container
     container.add_providers(*_CONNECTION_PROVIDERS)
 
-    def _enter_request() -> None:
-        connection = flask.request._get_current_object()  # noqa: SLF001  # ty: ignore[unresolved-attribute]
-        match = integrations.bind(flask_request_provider, connection)
-        child = container.build_child_container(scope=match.scope, context=match.context)
-        setattr(g, _CHILD_CONTAINER_ATTR, child)
-
     app.before_request(_enter_request)
     app.teardown_appcontext(_close_request)
+    # Flask runs these in reverse, so putting the close first makes it run last.
+    funcs = app.teardown_appcontext_funcs
+    funcs[:] = [_close_request, *(func for func in funcs if func is not _close_request)]
     if auto_inject:
         _inject_views(app)
     return container

tests/test_request_hooks.py

@@ -0,0 +1,191 @@
+import typing
+
+import flask
+import pytest
+from flask import Flask, Response, g, stream_with_context
+from modern_di import Container, Group, Scope, exceptions, providers
+
+from modern_di_flask import FromDI, inject, setup_di
+
+
+class Session:
+    pass
+
+
+def _make_group(events: list[str]) -> type[Group]:
+    class Deps(Group):
+        session = providers.Factory(
+            Session,
+            scope=Scope.REQUEST,
+            cache=providers.CacheSettings(finalizer=lambda _: events.append("finalized")),
+        )
+
+    return Deps
+
+
+@inject
+def _use(session: typing.Annotated[Session, FromDI(Session)]) -> Session:
+    return session
+
+
+def test_before_request_registered_before_setup_di_shares_the_request_container(app: Flask) -> None:
+    events: list[str] = []
+    seen: list[Session] = []
+
+    @app.before_request
+    def early() -> None:
+        seen.append(_use())
+
+    @app.route("/")
+    def index() -> str:
+        seen.append(_use())
+        return "ok"
+
+    with Container(groups=[_make_group(events)]) as root:
+        setup_di(app, root)
+        assert app.test_client().get("/").text == "ok"
+    assert len(seen) == 2  # noqa: PLR2004
+    assert seen[0] is seen[1]
+    assert events == ["finalized"]
+
+
+@pytest.mark.parametrize("order", ["before", "after"])
+def test_teardown_appcontext_handler_runs_before_the_close(app: Flask, order: str) -> None:
+    events: list[str] = []
+
+    def teardown(_exception: BaseException | None) -> None:
+        _use()
+        events.append("teardown")
+
+    app.add_url_rule("/", "index", lambda: (_use(), "ok")[1])
+    with Container(groups=[_make_group(events)]) as root:
+        if order == "before":
+            app.teardown_appcontext(teardown)
+        setup_di(app, root)
+        if order == "after":
+            app.teardown_appcontext(teardown)
+        assert app.test_client().get("/").text == "ok"
+    assert events == ["teardown", "finalized"]
+
+
+def test_teardown_request_after_short_circuit_builds_the_request_container(app: Flask) -> None:
+    events: list[str] = []
+
+    @app.before_request
+    def short_circuit() -> str:
+        return "short"
+
+    @app.teardown_request
+    def teardown(_exception: BaseException | None) -> None:
+        _use()
+        events.append("teardown")
+
+    with Container(groups=[_make_group(events)]) as root:
+        setup_di(app, root)
+        assert app.test_client().get("/").text == "short"
+    assert events == ["teardown", "finalized"]
+
+
+def test_teardown_does_not_build_a_request_container(app: Flask, container: Container) -> None:
+    built: list[bool] = []
+
+    @app.before_request
+    def short_circuit() -> str:
+        return "short"
+
+    def record(_sender: Flask, **_kwargs: object) -> None:
+        built.append(hasattr(g, "modern_di_request_container"))
+
+    setup_di(app, container)
+    with flask.appcontext_tearing_down.connected_to(record, app):
+        assert app.test_client().get("/").text == "short"
+    assert built == [False]
+
+
+def test_request_container_is_not_rebuilt_after_close(app: Flask) -> None:
+    events: list[str] = []
+    with Container(groups=[_make_group(events)]) as root:
+        setup_di(app, root)
+        with app.test_request_context("/"):
+            _use()
+            app.do_teardown_appcontext()
+            with pytest.raises(exceptions.ContainerClosedError):
+                _use()
+    assert events == ["finalized"]
+
+
+def test_streamed_response_never_rebuilds_the_request_container(app: Flask) -> None:
+    events: list[str] = []
+    outcomes: list[str] = []
+
+    @app.route("/")
+    def index() -> Response:
+        session = _use()
+
+        def generate() -> typing.Iterator[str]:
+            yield "a"
+            try:
+                outcomes.append("same" if _use() is session else "rebuilt")
+            except exceptions.ContainerClosedError:
+                outcomes.append("closed")
+
+        return Response(stream_with_context(generate()))
+
+    with Container(groups=[_make_group(events)]) as root:
+        setup_di(app, root)
+        assert app.test_client().get("/").text == "a"
+    # Flask 3.0 keeps the context for the whole stream; 3.1.2+ tears it down before iterating.
+    assert outcomes in (["same"], ["closed"])
+    assert events == ["finalized"]
+
+
+def test_requests_in_one_app_context_get_separate_request_containers(app: Flask) -> None:
+    events: list[str] = []
+    seen: list[Session] = []
+
+    @app.before_request
+    def early() -> None:
+        seen.append(_use())
+
+    app.add_url_rule("/", "index", lambda: "ok")
+    with Container(groups=[_make_group(events)]) as root:
+        setup_di(app, root)
+        with app.app_context():
+            client = app.test_client()
+            client.get("/")
+            client.get("/")
+    assert len(seen) == 2  # noqa: PLR2004
+    assert seen[0] is not seen[1]
+
+
+def test_setup_di_twice_registers_one_close_first(app: Flask, container: Container) -> None:
+    earlier = app.teardown_appcontext(lambda _exception: None)
+    setup_di(app, container)
+    with Container(groups=[_make_group([])]) as other:
+        setup_di(app, other)
+    close = app.teardown_appcontext_funcs[0]
+    assert app.teardown_appcontext_funcs == [close, earlier]
+
+
+def test_apps_resolve_from_their_own_containers() -> None:
+    first_events: list[str] = []
+    second_events: list[str] = []
+    first, second = Flask("first"), Flask("second")
+    for each in (first, second):
+        each.add_url_rule("/", "index", lambda: (_use(), "ok")[1])
+    with (
+        Container(groups=[_make_group(first_events)]) as first_root,
+        Container(groups=[_make_group(second_events)]) as second_root,
+    ):
+        setup_di(first, first_root)
+        setup_di(second, second_root)
+        first.test_client().get("/")
+        assert (first_events, second_events) == (["finalized"], [])
+        second.test_client().get("/")
+        assert (first_events, second_events) == (["finalized"], ["finalized"])
+
+
+def test_inject_outside_a_request_without_container_raises(app: Flask, container: Container) -> None:
+    setup_di(app, container)
+    with app.app_context(), pytest.raises(RuntimeError, match=r"setup_di\(app, container\)"):
+        _use()