From 468bd2aafbd7019ac851801988a2a8f64dda0f88 Mon Sep 17 00:00:00 2001 From: Artur Shiriev Date: Sat, 17 Oct 2026 15:55:16 +0300 Subject: [PATCH] feat: add manage_lifespan flag to setup_di --- CONTEXT.md | 4 +- README.md | 20 +++++++++- modern_di_taskiq/main.py | 4 +- tests/test_lifespan.py | 82 ++++++++++++++++++++++++++++++++++++++-- 4 files changed, 102 insertions(+), 8 deletions(-) diff --git a/CONTEXT.md b/CONTEXT.md index 4ff6f69..498833a 100644 --- a/CONTEXT.md +++ b/CONTEXT.md @@ -18,8 +18,8 @@ package's own. **Root container**: The `Container` that `setup_di` attaches to `broker.state`: the one a caller constructs and hands -in, the one `fetch_di_container` returns, and the one the worker's startup and shutdown events open -and close. +in, and the one `fetch_di_container` returns. By default the worker's startup and shutdown events +open and close it; with `setup_di(..., manage_lifespan=False)` the caller does. _Avoid_: APP-scope container — that names a scope where what matters is the position at the top of the tree, and `Scope.APP` is `modern-di`'s word for a band in the hierarchy, not for this object. diff --git a/README.md b/README.md index 5f807cd..b01241b 100644 --- a/README.md +++ b/README.md @@ -75,11 +75,29 @@ async def greet( The `WORKER_STARTUP`/`WORKER_SHUTDOWN` events fire when the broker's worker process starts and stops, so a script that calls tasks directly (like `InMemoryBroker` in a test) must drive the container lifecycle itself, for example with `async with broker: ...`. Providers resolve synchronously, but the per-task `Scope.REQUEST` child is closed with `close_async()`, so REQUEST-scoped providers can have async finalizers. `taskiq.TaskiqMessage` is resolvable within DI via the pre-built `taskiq_message_provider` context provider. +### Sharing a container + +When one container serves several apps, for example a taskiq worker and a web app in the same process, only one of them should open and close it. Pass `manage_lifespan=False` and own the root container yourself: + +```python +container = Container(groups=[AppGroup]) +setup_di(broker, container, manage_lifespan=False) + + +async def main() -> None: + async with container: + await broker.startup() + ... + await broker.shutdown() +``` + +The worker events then leave the root container alone. Each task still gets its own `Scope.REQUEST` child, closed when the task ends. APP-scoped finalizers run once, when the `async with` block exits. + ## API | Symbol | Description | |---|---| -| `setup_di(broker, container)` | Stores the root container on `broker.state`, opens/closes it on worker startup/shutdown, and registers `taskiq_message_provider`. Returns the container | +| `setup_di(broker, container, *, manage_lifespan=True)` | Stores the root container on `broker.state`, opens/closes it on worker startup/shutdown, and registers `taskiq_message_provider`. With `manage_lifespan=False` it never opens or closes the root container, so the caller owns it; per-task children are built and closed as before. Returns the container | | `FromDI(dependency, *, use_cache=True)` | taskiq dependency that resolves a provider instance or a type from the task's `Scope.REQUEST` child container, building that child on first use; use it as `Annotated[T, FromDI(...)]` in task signatures. `use_cache` is passed to `TaskiqDepends`. Raises `RuntimeError` naming `setup_di` when a task reaches it without `setup_di` called | | `fetch_di_container(broker)` | Returns the root container attached to the taskiq broker | | `taskiq_message_provider` | `ContextProvider` for the current `taskiq.TaskiqMessage` (`REQUEST` scope) | diff --git a/modern_di_taskiq/main.py b/modern_di_taskiq/main.py index 7f18bd3..2312a3e 100644 --- a/modern_di_taskiq/main.py +++ b/modern_di_taskiq/main.py @@ -16,9 +16,11 @@ _CONNECTION_PROVIDERS = (taskiq_message_provider,) -def setup_di(broker: AsyncBroker, container: Container) -> Container: +def setup_di(broker: AsyncBroker, container: Container, *, manage_lifespan: bool = True) -> Container: setattr(broker.state, _ROOT_CONTAINER_ATTR, container) container.add_providers(*_CONNECTION_PROVIDERS) + if not manage_lifespan: + return container def _open_root(_state: TaskiqState) -> None: container.open() diff --git a/tests/test_lifespan.py b/tests/test_lifespan.py index 90d7308..26fbd4e 100644 --- a/tests/test_lifespan.py +++ b/tests/test_lifespan.py @@ -1,9 +1,11 @@ -from modern_di import Container -from taskiq import InMemoryBroker +import typing + +from modern_di import Container, Group, Scope, providers +from taskiq import AsyncTaskiqDecoratedTask, InMemoryBroker import modern_di_taskiq -from modern_di_taskiq import fetch_di_container -from tests.dependencies import Dependencies +from modern_di_taskiq import FromDI, fetch_di_container, setup_di +from tests.dependencies import Dependencies, SimpleCreator def test_fetch_returns_the_same_container(broker: InMemoryBroker) -> None: @@ -40,3 +42,75 @@ async def test_restart_reopens_without_error(broker: InMemoryBroker) -> None: await broker.startup() # second cycle must not raise ContainerClosedError assert container.closed is False await broker.shutdown() + + +def _unmanaged_broker( + teardowns: list[str], +) -> tuple[InMemoryBroker, Container, AsyncTaskiqDecoratedTask[..., typing.Any]]: + class Resources(Group): + app_resource = providers.Factory( + scope=Scope.APP, + creator=SimpleCreator, + kwargs={"dep1": "app"}, + bound_type=None, + cache=providers.CacheSettings(finalizer=lambda _: teardowns.append("app")), + ) + task_resource = providers.Factory( + scope=Scope.REQUEST, + creator=SimpleCreator, + kwargs={"dep1": "task"}, + bound_type=None, + cache=providers.CacheSettings(finalizer=lambda _: teardowns.append("task")), + ) + + broker = InMemoryBroker() + container = setup_di(broker, Container(groups=[Resources]), manage_lifespan=False) + + @broker.task(task_name="use_resources") + async def use_resources( + app_resource: typing.Annotated[SimpleCreator, FromDI(Resources.app_resource)], + task_resource: typing.Annotated[SimpleCreator, FromDI(Resources.task_resource)], + ) -> str: + return f"{app_resource.dep1}+{task_resource.dep1}" + + return broker, container, use_resources + + +async def test_manage_lifespan_false_leaves_the_root_open_across_a_worker_cycle() -> None: + teardowns: list[str] = [] + broker, container, use_resources = _unmanaged_broker(teardowns) + + await broker.startup() + result = await (await use_resources.kiq()).wait_result() + await broker.shutdown() + + assert result.is_err is False + assert container.closed is False + assert "app" not in teardowns + + +async def test_manage_lifespan_false_still_builds_and_closes_the_per_task_child() -> None: + teardowns: list[str] = [] + broker, _container, use_resources = _unmanaged_broker(teardowns) + + await broker.startup() + result = await (await use_resources.kiq()).wait_result() + await broker.shutdown() + + assert result.return_value == "app+task" + assert teardowns == ["task"] + + +async def test_manage_lifespan_false_lets_the_caller_own_the_root() -> None: + teardowns: list[str] = [] + broker, container, use_resources = _unmanaged_broker(teardowns) + + async with container: + await broker.startup() + result = await (await use_resources.kiq()).wait_result() + await broker.shutdown() + assert teardowns == ["task"] + + assert result.is_err is False + assert container.closed is True + assert teardowns == ["task", "app"]