Support lifespan state

starlette/applications.py

@@ -9,7 +9,7 @@
 from starlette.requests import Request
 from starlette.responses import Response
 from starlette.routing import BaseRoute, Router
-from starlette.types import ASGIApp, Receive, Scope, Send
+from starlette.types import ASGIApp, Lifespan, Receive, Scope, Send
 
 
 class Starlette:
@@ -55,9 +55,7 @@ def __init__(
         ] = None,
         on_startup: typing.Optional[typing.Sequence[typing.Callable]] = None,
         on_shutdown: typing.Optional[typing.Sequence[typing.Callable]] = None,
-        lifespan: typing.Optional[
-            typing.Callable[["Starlette"], typing.AsyncContextManager]
-        ] = None,
+        lifespan: typing.Optional[Lifespan] = None,
     ) -> None:
         # The lifespan context function is a newer style that replaces
         # on_startup / on_shutdown handlers. Use one or the other, not both.

starlette/routing.py

@@ -17,7 +17,7 @@
 from starlette.middleware import Middleware
 from starlette.requests import Request
 from starlette.responses import PlainTextResponse, RedirectResponse
-from starlette.types import ASGIApp, Receive, Scope, Send
+from starlette.types import ASGIApp, Lifespan, Receive, Scope, Send, StatelessLifespan
 from starlette.websockets import WebSocket, WebSocketClose
 
 
@@ -558,17 +558,25 @@ def wrapper(app: typing.Any) -> _AsyncLiftContextManager:
     return wrapper
 
 
+_TDefaultLifespan = typing.TypeVar("_TDefaultLifespan", bound="_DefaultLifespan")
+
+
 class _DefaultLifespan:
     def __init__(self, router: "Router"):
         self._router = router
 
     async def __aenter__(self) -> None:
-        await self._router.startup()
+        await self._router.startup(state=self._state)
 
     async def __aexit__(self, *exc_info: object) -> None:
-        await self._router.shutdown()
-
-    def __call__(self: _T, app: object) -> _T:
+        await self._router.shutdown(state=self._state)
+
+    def __call__(
+        self: _TDefaultLifespan,
+        app: object,
+        state: typing.Optional[typing.Dict[str, typing.Any]],
+    ) -> _TDefaultLifespan:
+        self._state = state
         return self
 
 
@@ -580,9 +588,7 @@ def __init__(
         default: typing.Optional[ASGIApp] = None,
         on_startup: typing.Optional[typing.Sequence[typing.Callable]] = None,
         on_shutdown: typing.Optional[typing.Sequence[typing.Callable]] = None,
-        lifespan: typing.Optional[
-            typing.Callable[[typing.Any], typing.AsyncContextManager]
-        ] = None,
+        lifespan: typing.Optional[Lifespan] = None,
     ) -> None:
         self.routes = [] if routes is None else list(routes)
         self.redirect_slashes = redirect_slashes
@@ -591,10 +597,7 @@ def __init__(
         self.on_shutdown = [] if on_shutdown is None else list(on_shutdown)
 
         if lifespan is None:
-            self.lifespan_context: typing.Callable[
-                [typing.Any], typing.AsyncContextManager
-            ] = _DefaultLifespan(self)
-
+            self.lifespan_context: Lifespan = _DefaultLifespan(self)
         elif inspect.isasyncgenfunction(lifespan):
             warnings.warn(
                 "async generator function lifespans are deprecated, "
@@ -639,21 +642,31 @@ def url_path_for(self, name: str, **path_params: typing.Any) -> URLPath:
                 pass
         raise NoMatchFound(name, path_params)
 
-    async def startup(self) -> None:
+    async def startup(
+        self, state: typing.Optional[typing.Dict[str, typing.Any]]
+    ) -> None:
         """
         Run any `.on_startup` event handlers.
         """
         for handler in self.on_startup:
+            sig = inspect.signature(handler)
+            if len(sig.parameters) == 1 and state is not None:
+                handler = functools.partial(handler, state)
             if is_async_callable(handler):
                 await handler()
             else:
                 handler()
 
-    async def shutdown(self) -> None:
+    async def shutdown(
+        self, state: typing.Optional[typing.Dict[str, typing.Any]]
+    ) -> None:
         """
         Run any `.on_shutdown` event handlers.
         """
         for handler in self.on_shutdown:
+            sig = inspect.signature(handler)
+            if len(sig.parameters) == 1 and state is not None:
+                handler = functools.partial(handler, state)
             if is_async_callable(handler):
                 await handler()
             else:
@@ -666,9 +679,18 @@ async def lifespan(self, scope: Scope, receive: Receive, send: Send) -> None:
         """
         started = False
         app = scope.get("app")
+        state = scope.get("state")
         await receive()
         try:
-            async with self.lifespan_context(app):
+            lifespan_context: Lifespan
+            if (
+                len(inspect.signature(self.lifespan_context).parameters) == 2
+                and state is not None
+            ):
+                lifespan_context = functools.partial(self.lifespan_context, state=state)
+            else:
+                lifespan_context = typing.cast(StatelessLifespan, self.lifespan_context)
+            async with lifespan_context(app):
                 await send({"type": "lifespan.startup.complete"})
                 started = True
                 await receive()

starlette/testclient.py

@@ -188,11 +188,14 @@ def __init__(
         portal_factory: _PortalFactoryType,
         raise_server_exceptions: bool = True,
         root_path: str = "",
+        *,
+        app_state: typing.Dict[str, typing.Any],
     ) -> None:
         self.app = app
         self.raise_server_exceptions = raise_server_exceptions
         self.root_path = root_path
         self.portal_factory = portal_factory
+        self.app_state = app_state
 
     def handle_request(self, request: httpx.Request) -> httpx.Response:
         scheme = request.url.scheme
@@ -243,6 +246,7 @@ def handle_request(self, request: httpx.Request) -> httpx.Response:
                 "client": ["testclient", 50000],
                 "server": [host, port],
                 "subprotocols": subprotocols,
+                "state": self.app_state.copy(),
             }
             session = WebSocketTestSession(self.app, scope, self.portal_factory)
             raise _Upgrade(session)
@@ -260,6 +264,7 @@ def handle_request(self, request: httpx.Request) -> httpx.Response:
             "client": ["testclient", 50000],
             "server": [host, port],
             "extensions": {"http.response.debug": {}},
+            "state": self.app_state.copy(),
         }
 
         request_complete = False
@@ -380,11 +385,13 @@ def __init__(
             app = typing.cast(ASGI2App, app)  # type: ignore[assignment]
             asgi_app = _WrapASGI2(app)  # type: ignore[arg-type]
         self.app = asgi_app
+        self.app_state: typing.Dict[str, typing.Any] = {}
         transport = _TestClientTransport(
             self.app,
             portal_factory=self._portal_factory,
             raise_server_exceptions=raise_server_exceptions,
             root_path=root_path,
+            app_state=self.app_state,
         )
         if headers is None:
             headers = {}
@@ -749,7 +756,7 @@ def __exit__(self, *args: typing.Any) -> None:
         self.exit_stack.close()
 
     async def lifespan(self) -> None:
-        scope = {"type": "lifespan"}
+        scope = {"type": "lifespan", "state": self.app_state}
         try:
             await self.app(scope, self.stream_receive.receive, self.stream_send.send)
         finally:

starlette/types.py

@@ -7,3 +7,9 @@
 Send = typing.Callable[[Message], typing.Awaitable[None]]
 
 ASGIApp = typing.Callable[[Scope, Receive, Send], typing.Awaitable[None]]
+
+StatelessLifespan = typing.Callable[[object], typing.AsyncContextManager]
+StateLifespan = typing.Callable[
+    [typing.Any, typing.Dict[str, typing.Any]], typing.AsyncContextManager
+]
+Lifespan = typing.Union[StatelessLifespan, StateLifespan]

tests/test_routing.py

@@ -669,6 +669,53 @@ def run_shutdown():
     assert shutdown_complete
 
 
+def test_lifespan_with_state(test_client_factory):
+    startup_complete = False
+    shutdown_complete = False
+
+    async def hello_world(request):
+        # modifications to the state should not leak across requests
+        assert request.state.count == 0
+        # modify the state, this should not leak to the lifespan or other requests
+        request.state.count += 1
+        # since state.list is a mutable object this modification _will_ leak across
+        # requests and to the lifespan
+        request.state.list.append(1)
+        return PlainTextResponse("hello, world")
+
+    async def run_startup(state):
+        nonlocal startup_complete
+        startup_complete = True
+        state["count"] = 0
+        state["list"] = []
+
+    async def run_shutdown(state):
+        nonlocal shutdown_complete
+        shutdown_complete = True
+        # modifications made to the state from a request do not leak to the lifespan
+        assert state["count"] == 0
+        # unless of course the request mutates a mutable object that is referenced
+        # via state
+        assert state["list"] == [1, 1]
+
+    app = Router(
+        on_startup=[run_startup],
+        on_shutdown=[run_shutdown],
+        routes=[Route("/", hello_world)],
+    )
+
+    assert not startup_complete
+    assert not shutdown_complete
+    with test_client_factory(app) as client:
+        assert startup_complete
+        assert not shutdown_complete
+        client.get("/")
+        # Calling it a second time to ensure that the state is preserved.
+        client.get("/")
+    assert startup_complete
+    assert shutdown_complete
+
+
 def test_raise_on_startup(test_client_factory):
     def run_startup():
         raise RuntimeError()