feat: add manage_lifespan flag to setup_di

CONTEXT.md

@@ -17,7 +17,8 @@ one. Celery's vocabulary is Celery's — task, worker, broker, beat, and the poo
 
 **Root container**:
 The `Scope.APP` `Container` handed to `setup_di` and registered for the Celery app, outside its
-`conf`; the worker lifecycle signals open and close it, and every per-task child descends from it.
+`conf`; the worker lifecycle signals open and close it unless `setup_di` gets `manage_lifespan=False`,
+which leaves that to the user, and every per-task child descends from it.
 `fetch_di_container` returns this one.
 _Avoid_: APP container, APP-scope container — both have been used for it, and neither says the thing
 that matters at a call site, which is that this is the container a per-task child is built *from*.

README.md

@@ -80,11 +80,29 @@ These worker signals fire only when a real `celery worker` starts and stops, so
 
 Celery runs its prefork pool on billiard, its own copy of multiprocessing, which starts children with the fork method by default on macOS and Linux. With spawn (Windows, or `billiard.set_start_method("spawn", force=True)`), each child unpickles the app and re-imports your task modules, so `setup_di` runs again in the child. The `worker_process_init` handler then registers the child's container for the app the child runs. It matches apps by `main` name, so give each app in a process its own name. Under spawn, Celery's prefork pool also needs `FORKED_BY_MULTIPROCESSING=1` in the environment. Without it tasks fail inside Celery's tracer, with or without modern-di.
 
+### Sharing a container
+
+To share one container between the Celery app and other code in the same process, such as another app built on modern-di, pass `manage_lifespan=False` and open and close the root container yourself. The worker signals then leave it alone, and each task call still gets its own child container:
+
+```python
+from modern_di_celery import setup_di
+
+
+setup_di(app, container, manage_lifespan=False)
+
+
+async def main() -> None:
+    async with container:
+        ...  # run the apps that share the container
+```
+
+Under the prefork pool each forked child inherits the open container. With `manage_lifespan=False` nothing closes it in the child, so close it yourself if its `Scope.APP` finalizers must run there, for example from your own `worker_process_shutdown` handler that calls `container.close_sync()`.
+
 ## API
 
 | Symbol | Description |
 |---|---|
-| `setup_di(app, container)` | Registers the root container (`Scope.APP`) for the app, outside `app.conf`, and opens/closes it on `worker_process_init`/`worker_process_shutdown` and `worker_init`/`worker_shutdown`. Returns the container |
+| `setup_di(app, container, *, manage_lifespan=True)` | Registers the root container (`Scope.APP`) for the app, outside `app.conf`, and opens/closes it on `worker_process_init`/`worker_process_shutdown` and `worker_init`/`worker_shutdown`. With `manage_lifespan=False` it never opens or closes the root container, and `@inject` still builds and closes a child container per task call. Returns the container |
 | `FromDI(dependency)` | Inert marker for `Annotated[T, FromDI(...)]` in task signatures; accepts a provider instance or a type |
 | `inject(task)` | Decorator that builds a `Scope.REQUEST` child per call, resolves the `FromDI`-annotated parameters, and closes the child with `close_sync()` afterwards. Raises `TypeError` at decoration if the task also declares `*args`/`**kwargs`. Raises `RuntimeError` naming `setup_di` when a task reaches it without `setup_di` called |
 | `DITask` | `Task` subclass that applies `@inject` to a task's `run` automatically; pass `task_cls=DITask` to `Celery(...)` or `base=DITask` to `@app.task(...)` |

modern_di_celery/main.py

@@ -33,28 +33,32 @@ def _bind_current_app(main: str | None, container: Container) -> Celery:
     return process_app
 
 
-def setup_di(app: Celery, container: Container) -> Container:
+def setup_di(app: Celery, container: Container, *, manage_lifespan: bool = True) -> Container:
     _ROOT_CONTAINERS[app] = container
     main = app.main
 
     def _owns(owner: Celery) -> bool:
         return _ROOT_CONTAINERS.get(owner) is container
 
-    def _open_for_worker(sender: object = None, **_: typing.Any) -> None:  # noqa: ANN401
-        if _owns(_worker_app(sender)):
+    def _open_for(owner: Celery) -> None:
+        if manage_lifespan and _owns(owner):
             container.open()
 
-    def _close_for_worker(sender: object = None, **_: typing.Any) -> None:  # noqa: ANN401
-        if _owns(_worker_app(sender)):
+    def _close_for(owner: Celery) -> None:
+        if manage_lifespan and _owns(owner):
             container.close_sync()
 
+    def _open_for_worker(sender: object = None, **_: typing.Any) -> None:  # noqa: ANN401
+        _open_for(_worker_app(sender))
+
+    def _close_for_worker(sender: object = None, **_: typing.Any) -> None:  # noqa: ANN401
+        _close_for(_worker_app(sender))
+
     def _open_for_process(**_: typing.Any) -> None:  # noqa: ANN401
-        if _owns(_bind_current_app(main, container)):
-            container.open()
+        _open_for(_bind_current_app(main, container))
 
     def _close_for_process(**_: typing.Any) -> None:  # noqa: ANN401
-        if _owns(_current_app()):
-            container.close_sync()
+        _close_for(_current_app())
 
     # weak=False is required: Celery signals default to weak refs, which would
     # garbage-collect these handlers before a worker process ever fires them.

tests/test_lifespan.py

@@ -1,7 +1,8 @@
+import asyncio
 import typing
 
 from celery import Celery, signals
-from modern_di import Container
+from modern_di import Container, Group, Scope, providers
 
 import modern_di_celery
 from modern_di_celery import FromDI, fetch_di_container, inject
@@ -103,3 +104,91 @@ def test_process_signals_act_only_on_the_current_app() -> None:
     container_b.close_sync()
     signals.worker_process_init.send(sender=None)
     assert (container_a.closed, container_b.closed) == (False, True)
+
+
+class _AppResource:
+    pass
+
+
+class _TaskResource:
+    pass
+
+
+def _tracked_container(finalized: list[str]) -> Container:
+    class Tracked(Group):
+        app_resource = providers.Factory(
+            _AppResource,
+            scope=Scope.APP,
+            cache=providers.CacheSettings(finalizer=lambda _: finalized.append("app")),
+        )
+        task_resource = providers.Factory(
+            _TaskResource,
+            scope=Scope.REQUEST,
+            cache=providers.CacheSettings(finalizer=lambda _: finalized.append("task")),
+        )
+
+    return Container(groups=[Tracked])
+
+
+def _unmanaged_app(main: str, container: Container) -> Celery:
+    app = Celery(main, broker="memory://", backend="cache+memory://")
+    app.conf.task_always_eager = True
+    app.conf.task_store_eager_result = True
+    modern_di_celery.setup_di(app, container, manage_lifespan=False)
+
+    @app.task(name="uses_both_scopes")
+    @inject
+    def uses_both_scopes(
+        app_resource: typing.Annotated[_AppResource, FromDI(_AppResource)],
+        task_resource: typing.Annotated[_TaskResource, FromDI(_TaskResource)],
+    ) -> bool:
+        return isinstance(app_resource, _AppResource) and isinstance(task_resource, _TaskResource)
+
+    return app
+
+
+def _run_worker_lifecycle(app: Celery) -> bool:
+    worker = _FakeWorker(app)
+    signals.worker_init.send(sender=worker)
+    signals.worker_process_init.send(sender=None)
+    result = app.tasks["uses_both_scopes"].delay().get()
+    signals.worker_process_shutdown.send(sender=None)
+    signals.worker_shutdown.send(sender=worker)
+    return result
+
+
+def test_manage_lifespan_false_leaves_the_root_open_and_unfinalized() -> None:
+    finalized: list[str] = []
+    container = _tracked_container(finalized)
+    app = _unmanaged_app("unmanaged-lifecycle", container)
+
+    assert _run_worker_lifecycle(app) is True
+    assert container.closed is False
+    assert finalized == ["task"]
+
+
+def test_manage_lifespan_false_never_opens_the_root() -> None:
+    container = _tracked_container([])
+    app = _unmanaged_app("unmanaged-closed", container)
+    container.close_sync()
+
+    signals.worker_init.send(sender=_FakeWorker(app))
+    signals.worker_process_init.send(sender=None)
+
+    assert container.closed is True
+
+
+def test_manage_lifespan_false_lets_the_user_own_the_root() -> None:
+    finalized: list[str] = []
+    container = _tracked_container(finalized)
+    app = _unmanaged_app("unmanaged-owned", container)
+
+    async def run() -> bool:
+        async with container:
+            result = _run_worker_lifecycle(app)
+            assert finalized == ["task"]
+        return result
+
+    assert asyncio.run(run()) is True
+    assert container.closed is True
+    assert finalized == ["task", "app"]

tests/test_spawn.py

@@ -63,3 +63,17 @@ def test_an_app_with_a_container_keeps_it_when_a_later_handler_fires() -> None:
     signals.worker_process_init.send(sender=None)
 
     assert fetch_di_container(current) is current_container
+
+
+def test_a_spawn_child_binds_its_container_without_opening_it_when_lifespan_is_unmanaged() -> None:
+    payload = pickle.dumps(_make_app("spawn-unmanaged", set_as_current=False))
+    unpickled = pickle.loads(payload)  # noqa: S301
+    recreated = _make_app("spawn-unmanaged", set_as_current=False)
+    container = setup_di(recreated, Container(groups=[Dependencies]), manage_lifespan=False)
+    container.close_sync()
+    unpickled.set_current()
+
+    signals.worker_process_init.send(sender=None)
+
+    assert fetch_di_container(typing.cast(Celery, current_app)) is container
+    assert container.closed is True