From 121a093facc8e0e761284dbc0be75428018c199d Mon Sep 17 00:00:00 2001 From: Artur Shiriev Date: Sat, 17 Oct 2026 15:53:14 +0300 Subject: [PATCH] feat: add manage_lifespan to setup_di --- CONTEXT.md | 3 +- README.md | 18 +++++++++++- modern_di_faststream/main.py | 9 ++++-- tests/dependencies.py | 20 +++++++++++++ tests/test_lifespan.py | 56 ++++++++++++++++++++++++++++++++++-- 5 files changed, 99 insertions(+), 7 deletions(-) diff --git a/CONTEXT.md b/CONTEXT.md index ed56e90..dc17b96 100644 --- a/CONTEXT.md +++ b/CONTEXT.md @@ -17,7 +17,8 @@ uses each in FastStream's sense. The three below are this package's own. **Root container**: The `Container` handed to `setup_di` and stored on the app's `ContextRepo`; `fetch_di_container` -reads it back. Its lifetime tracks the app's, reopened on startup and closed after shutdown. +reads it back. By default its lifetime tracks the app's, reopened on startup and closed after shutdown; +with `manage_lifespan=False` the user opens and closes it instead. _Avoid_: app container, app-scoped container — accurate but unnamed; the code calls it the root container, and one of the two names has to win. diff --git a/README.md b/README.md index 336458f..dbe4c52 100644 --- a/README.md +++ b/README.md @@ -71,11 +71,27 @@ async def handle(name: str, handler: GreetingHandler = FromDI(Dependencies.handl The current `StreamMessage` is resolvable within DI via the pre-built `faststream_message_provider` context provider. +### Sharing a container + +By default the app opens the container on startup and closes it after shutdown, which runs every `APP`-scoped finalizer. When the container outlives the app, or other apps share it, pass `manage_lifespan=False` and open and close it yourself: + +```python +container = Container(groups=[Dependencies]) +setup_di(app, container, manage_lifespan=False) + + +async def main() -> None: + async with container: + await app.run() +``` + +When several apps share one container, exactly one of them should own its lifespan, or the first app to stop closes the container for the rest. + ## API | Symbol | Description | |---|---| -| `setup_di(app, container)` | Stores the container in the app context and registers startup/shutdown lifecycle hooks. On startup it reopens the container and adds the DI middleware to every broker of the app, including one added via `app.add_broker` after `setup_di` or by an `on_startup` hook registered before `setup_di`. After shutdown it closes the container. Raises `RuntimeError` at startup if the app has no broker by then. Returns the container | +| `setup_di(app, container, *, manage_lifespan=True)` | Stores the container in the app context and registers startup/shutdown lifecycle hooks. On startup it reopens the container and adds the DI middleware to every broker of the app, including one added via `app.add_broker` after `setup_di` or by an `on_startup` hook registered before `setup_di`. After shutdown it closes the container. With `manage_lifespan=False` it never opens or closes the container, and the middleware still builds and closes a request container per message. Raises `RuntimeError` at startup if the app has no broker by then. Returns the container | | `FromDI(dependency, *, use_cache=True, cast=False)` | FastStream `Depends` that resolves a provider (or type) from the request container. Raises `RuntimeError` naming `setup_di` when a message reaches it without the middleware installed | | `fetch_di_container(app)` | Returns the root container from the app context | | `faststream_message_provider` | `ContextProvider` for the current `faststream.StreamMessage` (`REQUEST` scope) | diff --git a/modern_di_faststream/main.py b/modern_di_faststream/main.py index f6781dd..fa475a6 100644 --- a/modern_di_faststream/main.py +++ b/modern_di_faststream/main.py @@ -70,6 +70,8 @@ def fetch_di_container(app_: faststream.FastStream | AsgiFastStream) -> Containe def setup_di( app: faststream.FastStream | AsgiFastStream, container: Container, + *, + manage_lifespan: bool = True, ) -> Container: container.add_providers(faststream_message_provider) app.context.set_global(_ROOT_CONTAINER_KEY, container) @@ -84,13 +86,14 @@ def install_middleware() -> None: if middleware_factory not in broker.config.broker_middlewares: broker.add_middleware(middleware_factory) + app.on_startup(install_middleware) # FastStream's lifecycle is callback-based, so the root container can't be # wrapped in ``async with``. Reopen it on startup (before the broker consumes) # to pair with the shutdown close, so a broker restart works instead of # raising ContainerClosedError. Reopening an already-open container is a no-op. - app.on_startup(container.open) - app.on_startup(install_middleware) - app.after_shutdown(container.close_async) + if manage_lifespan: + app.on_startup(container.open) + app.after_shutdown(container.close_async) return container diff --git a/tests/dependencies.py b/tests/dependencies.py index 0a78d17..7c807cb 100644 --- a/tests/dependencies.py +++ b/tests/dependencies.py @@ -15,6 +15,14 @@ class DependentCreator: dep1: SimpleCreator +@dataclasses.dataclass(kw_only=True, slots=True) +class Finalizable: + finalized: int = 0 + + def finalize(self) -> None: + self.finalized += 1 + + def fetch_message_is_processed_from_request(message: StreamMessage[typing.Any] | None = None) -> bool: return message.processed if message else False @@ -24,3 +32,15 @@ class Dependencies(Group): request_factory = providers.Factory(scope=Scope.REQUEST, creator=DependentCreator, bound_type=None) action_factory = providers.Factory(scope=Scope.ACTION, creator=DependentCreator, bound_type=None) message_is_processed = providers.Factory(scope=Scope.REQUEST, creator=fetch_message_is_processed_from_request) + app_finalizable = providers.Factory( + scope=Scope.APP, + creator=Finalizable, + bound_type=None, + cache=providers.CacheSettings(finalizer=Finalizable.finalize), + ) + request_finalizable = providers.Factory( + scope=Scope.REQUEST, + creator=Finalizable, + bound_type=None, + cache=providers.CacheSettings(finalizer=Finalizable.finalize), + ) diff --git a/tests/test_lifespan.py b/tests/test_lifespan.py index c71b0b2..891d59f 100644 --- a/tests/test_lifespan.py +++ b/tests/test_lifespan.py @@ -3,9 +3,10 @@ import faststream from faststream import TestApp from faststream.nats import NatsBroker, TestNatsBroker +from modern_di import Container -from modern_di_faststream import FromDI, fetch_di_container -from tests.dependencies import Dependencies, SimpleCreator +from modern_di_faststream import FromDI, fetch_di_container, setup_di +from tests.dependencies import Dependencies, Finalizable, SimpleCreator TEST_SUBJECT = "test" @@ -43,3 +44,54 @@ async def index_subscriber( async with TestApp(app): assert not container.closed await br.publish("test", TEST_SUBJECT) + + +def _unmanaged_app() -> tuple[faststream.FastStream, NatsBroker, Container, list[Finalizable]]: + broker = NatsBroker() + app = faststream.FastStream(broker) + container = setup_di(app, Container(groups=[Dependencies]), manage_lifespan=False) + request_instances: list[Finalizable] = [] + + @broker.subscriber(TEST_SUBJECT) + async def index_subscriber( + message: str, + app_instance: typing.Annotated[Finalizable, FromDI(Dependencies.app_finalizable)], + request_instance: typing.Annotated[Finalizable, FromDI(Dependencies.request_finalizable)], + ) -> None: + assert message == "test" + assert app_instance is container.resolve_provider(Dependencies.app_finalizable) + request_instances.append(request_instance) + + return app, broker, container, request_instances + + +async def test_manage_lifespan_false_leaves_the_root_open_after_shutdown() -> None: + app, broker, container, _ = _unmanaged_app() + + async with TestNatsBroker(broker) as br, TestApp(app): + await br.publish("test", TEST_SUBJECT) + + assert not container.closed + assert container.resolve_provider(Dependencies.app_finalizable).finalized == 0 + + +async def test_manage_lifespan_false_still_builds_and_closes_the_request_container() -> None: + app, broker, _, request_instances = _unmanaged_app() + + async with TestNatsBroker(broker) as br, TestApp(app): + await br.publish("test", TEST_SUBJECT) + + assert [instance.finalized for instance in request_instances] == [1] + + +async def test_manage_lifespan_false_lets_the_user_own_the_root() -> None: + app, broker, container, _ = _unmanaged_app() + + async with container: + async with TestNatsBroker(broker) as br, TestApp(app): + await br.publish("test", TEST_SUBJECT) + app_instance = container.resolve_provider(Dependencies.app_finalizable) + assert app_instance.finalized == 0 + + assert container.closed + assert app_instance.finalized == 1