Repository: uriyyo/fastapi-lifespan-manager (language: python)

Issue:
How to have one piece of state depend on another?
Say I want to have one piece of state that is an httpx client and another that is an OpenAI API client, which takes an httpx client on initialization. If I do

```py
manager = LifespanManager()

@manager.add
async def set_up_httpx(app: FastAPI) -> AsyncIterator[State]:
    async with httpx.AsyncClient() as client:
        yield {"http": client}

@manager.add
async def set_up_openapi(app: FastAPI) -> AsyncIterator[State]:
    http: httpx.AsyncClient = app.state.http
    client = AsyncOpenAI(http_client=http)
    yield {"openai": client}
```

the latter call does not receive the state from the earlier call because the state is not attached to the app here: https://github.com/uriyyo/fastapi-lifespan-manager/blob/808e65df3cfdce0fbe200fd8e48383669535a626/fastapi_lifespan_manager/lifespan_manager.py#L98 

Is there a recommended way to do this?


Discussion / hints:

Fix patch:
diff --git a/fastapi_lifespan_manager/lifespan_manager.py b/fastapi_lifespan_manager/lifespan_manager.py
index 81e9e04..277d70f 100644
--- a/fastapi_lifespan_manager/lifespan_manager.py
+++ b/fastapi_lifespan_manager/lifespan_manager.py
@@ -1,8 +1,9 @@
 from __future__ import annotations
 
 import inspect
-from contextlib import AsyncExitStack, asynccontextmanager, contextmanager
+from contextlib import AsyncExitStack, asynccontextmanager, contextmanager, suppress
 from dataclasses import dataclass
+from types import MappingProxyType
 from typing import (
     TYPE_CHECKING,
     Any,
@@ -25,37 +26,45 @@
 from fastapi.concurrency import contextmanager_in_threadpool
 from typing_extensions import Self
 
-from .types import AnyContextManager, AnyState, RawLifespan, TApp
+from .types import AnyContextManager, AnyState, RawLifespan, State, TApp
 
 
-def _maybe_call(call: Callable[..., AnyContextManager], app: Optional[TApp]) -> AnyContextManager:
-    try:
+def _maybe_call(call: Callable[..., AnyContextManager], app: Optional[TApp], state: State) -> AnyContextManager:
+    with suppress(TypeError):
+        return call(app, state)
+
+    with suppress(TypeError):
         return call(app)
-    except TypeError:
-        return call()
+
+    return call()
 
 
 def _convert_raw_lifespan_to_ctx(
     app: TApp,
+    state: State,
     raw_lifespan: RawLifespan[TApp],
 ) -> Union[ContextManager[AnyState], AsyncContextManager[AnyState]]:
     if inspect.isasyncgenfunction(raw_lifespan):
-        return _maybe_call(asynccontextmanager(raw_lifespan), app)
+        return _maybe_call(asynccontextmanager(raw_lifespan), app, state)
     if inspect.isgeneratorfunction(raw_lifespan):
-        return _maybe_call(contextmanager(raw_lifespan), app)
+        return _maybe_call(contextmanager(raw_lifespan), app, state)
 
-    return _maybe_call(cast(Callable[..., AnyContextManager], raw_lifespan), app)
+    return _maybe_call(cast(Callable[..., AnyContextManager], raw_lifespan), app, state)
 
 
 @asynccontextmanager
-async def _run_raw_lifespan(raw_lifespan: RawLifespan[TApp], app: TApp) -> AsyncIterator[AnyState]:
-    ctx = _convert_raw_lifespan_to_ctx(app, raw_lifespan)
+async def _run_raw_lifespan(
+    raw_lifespan: RawLifespan[TApp],
+    app: TApp,
+    state: State,
+) -> AsyncIterator[AnyState]:
+    ctx = _convert_raw_lifespan_to_ctx(app, state, raw_lifespan)
 
     actx: AsyncContextManager[AnyState]
     actx = contextmanager_in_threadpool(ctx) if isinstance(ctx, ContextManager) else ctx
 
-    async with actx as state:
-        yield state
+    async with actx as res:
+        yield res
 
 
 class LifespanManager(Generic[TApp]):
@@ -92,14 +101,13 @@ def include(self, other: Self) -> None:
     @asynccontextmanager
     async def __call__(self, app: TApp) -> AsyncIterator[AnyState]:
         async with AsyncExitStack() as astack:
-            state: Optional[Dict[str, Any]] = None
+            state: Dict[str, Any] = {}
+            state_proxy: State = MappingProxyType(state)
 
             for raw_lifespan in self.lifespans:
-                sub_state = await astack.enter_async_context(_run_raw_lifespan(raw_lifespan, app))
+                sub_state = await astack.enter_async_context(_run_raw_lifespan(raw_lifespan, app, state_proxy))
 
                 if sub_state:
-                    if state is None:
-                        state = {}
                     state.update(sub_state)
 
             yield state
diff --git a/fastapi_lifespan_manager/types.py b/fastapi_lifespan_manager/types.py
index f294dbb..e725674 100644
--- a/fastapi_lifespan_manager/types.py
+++ b/fastapi_lifespan_manager/types.py
@@ -13,6 +13,8 @@
 AnyIterator: TypeAlias = Union[Iterator[AnyState], AsyncIterator[AnyState]]
 
 RawLifespan: TypeAlias = Union[
+    # (app, state) -> ...
+    Callable[[TApp, State], Union[AnyContextManager, AnyIterator]],
     # (app) -> ...
     Callable[[TApp], Union[AnyContextManager, AnyIterator]],
     # () -> ...
diff --git a/pyproject.toml b/pyproject.toml
index ea2cb0b..2053fb2 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -50,6 +50,9 @@ exclude = '''
 
 [tool.ruff]
 line-length = 120
+target-version = "py38"
+
+[tool.ruff.lint]
 select = [
     "E", # pycodestyle errors
     "W", # pycodestyle warnings
@@ -100,9 +103,8 @@ exclude = [
 ]
 per-file-ignores = { }
 dummy-variable-rgx = "^(_+|(_+[a-zA-Z0-9_]*[a-zA-Z0-9]+?))$"
-target-version = "py38"
 
-[tool.ruff.mccabe]
+[tool.ruff.lint.mccabe]
 max-complexity = 10
 
 [tool.mypy]


Test patch:
diff --git a/tests/conftest.py b/tests/conftest.py
index 2347ceb..e69de29 100644
--- a/tests/conftest.py
+++ b/tests/conftest.py
@@ -1,10 +0,0 @@
-from asyncio import new_event_loop
-
-import pytest
-
-
-@pytest.fixture(scope="session")
-def event_loop():
-    loop = new_event_loop()
-    yield loop
-    loop.close()
diff --git a/tests/test_lifespan_manager.py b/tests/test_lifespan_manager.py
index c6277eb..0bd2c1e 100644
--- a/tests/test_lifespan_manager.py
+++ b/tests/test_lifespan_manager.py
@@ -137,6 +137,23 @@ async def _lifespan(_):
     assert called
 
 
+@pytest.mark.asyncio()
+async def test_state_inside_lifespan():
+    manager = LifespanManager()
+
+    @manager.add
+    async def _lifespan_1(app, state):
+        assert state == {}
+        yield {"foo": "bar"}
+
+    @manager.add
+    async def _lifespan_2(app, state):
+        assert state == {"foo": "bar"}
+        yield
+
+    await get_state(manager)
+
+
 @pytest.mark.asyncio()
 async def test_remove():
     manager = LifespanManager()
diff --git a/tests/test_mypy.py b/tests/test_mypy.py
index 05ec2f9..4e07a23 100644
--- a/tests/test_mypy.py
+++ b/tests/test_mypy.py
@@ -73,40 +73,39 @@ def _lifespan(app: Starlette) -> Iterator[NoState]:
 
 
 @pytest.mark.parametrize(
-    "no_app",
-    [False, True],
-    ids=["with-app", "no-app"],
+    "type_",
+    ["with-app", "with-state", "no-args"],
 )
 @pytest.mark.parametrize(
     "code",
     [
         """\
-    def _lifespan(app: FastAPI) -> Iterator[NoState]:
+    def _lifespan() -> Iterator[NoState]:
         yield""",
         """\
-    def _lifespan(app: FastAPI) -> Iterator[State]:
+    def _lifespan() -> Iterator[State]:
         yield {}""",
         """\
     @contextmanager
-    def _lifespan(app: FastAPI) -> Iterator[NoState]:
+    def _lifespan() -> Iterator[NoState]:
         yield""",
         """\
     @contextmanager
-    def _lifespan(app: FastAPI) -> Iterator[State]:
+    def _lifespan() -> Iterator[State]:
         yield {}""",
         """\
-    async def _lifespan(app: FastAPI) -> AsyncIterator[NoState]:
+    async def _lifespan() -> AsyncIterator[NoState]:
         yield""",
         """\
-    async def _lifespan(app: FastAPI) -> AsyncIterator[State]:
+    async def _lifespan() -> AsyncIterator[State]:
         yield {}""",
         """\
     @asynccontextmanager
-    async def _lifespan(app: FastAPI) -> AsyncIterator[NoState]:
+    async def _lifespan() -> AsyncIterator[NoState]:
         yield""",
         """\
     @asynccontextmanager
-    async def _lifespan(app: FastAPI) -> AsyncIterator[State]:
+    async def _lifespan() -> AsyncIterator[State]:
         yield {}""",
     ],
     ids=[
@@ -120,9 +119,11 @@ async def _lifespan(app: FastAPI) -> AsyncIterator[State]:
         "async-ctx-with-state",
     ],
 )
-def test_different_kind_of_lifespans(no_app, code):
-    if no_app:
-        code = code.replace("app: FastAPI", "")
+def test_different_kind_of_lifespans(type_, code):
+    if type_ == "with-app":
+        code = code.replace("()", "(app: FastAPI)")
+    if type_ == "with-state":
+        code = code.replace("()", "(app: FastAPI, state: State)")
 
     output = _run_mypy(
         f"""