feat: add manage_lifespan flag to setup_di

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.
 

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) |

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()

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"]