upload

app/components/lifespan.py

@@ -1,25 +1,59 @@
-from typing import TypeGuard, override
+from collections.abc import AsyncGenerator, Callable
+from typing import Any, TypeGuard, override
 
-from app.types_ import AnyScope, LifespanScope, Receive, ReceiveLifespan, Send
+from app.subroutines.asyncutils import agzip
+from app.types_ import (
+    AnyScope,
+    AsyncCallable,
+    LifespanScope,
+    Receive,
+    ReceiveLifespan,
+    Send,
+)
 
 from .base import Component as _Component
 
 
 class LifespanComponent(_Component[LifespanScope, ReceiveLifespan]):
+    startups: list[AsyncCallable[[], None]]
+    shutdowns: list[AsyncCallable[[], None]]
+    contexts: list[Callable[[], AsyncGenerator[Any, None]]]
+
+    def __init__(self, *args: Any, **kwds: Any) -> None:
+        super().__init__(*args, **kwds)
+        self.startups = []
+        self.shutdowns = []
+        self.contexts = []
+
     @override
     async def condition(self, scope: AnyScope) -> TypeGuard[LifespanScope]:
         return scope["type"] == "lifespan"
 
     @override
-    async def handle(self, scope: LifespanScope, receive: Receive[ReceiveLifespan], send: Send) -> None:
-        while True:
+    async def handle(
+        self, scope: LifespanScope, receive: Receive[ReceiveLifespan], send: Send
+    ) -> None:
+        message = await receive()
+        async for _ in agzip(*[ctx() for ctx in self.contexts]):
+            if message["type"] == "lifespan.startup":
+                for fn in self.startups:
+                    await fn()
+                await send({"type": "lifespan.startup.complete"})
+            elif message["type"] == "lifespan.shutdown":
+                for fn in self.shutdowns:
+                    await fn()
+                await send({"type": "lifespan.shutdown.complete"})
+                return
             message = await receive()
-            if message['type'] == 'lifespan.startup':
-                ... # Do some startup here!
-                print("Startup...")
-                await send({'type': 'lifespan.startup.complete'})
-            elif message['type'] == 'lifespan.shutdown':
-                ... # Do some shutdown here!
-                print("Shutdown...")
-                await send({'type': 'lifespan.shutdown.complete'})
-                return
+
+    def on_startup[Call_T: AsyncCallable[[], None]](self, fn: Call_T) -> Call_T:
+        self.startups.append(fn)
+        return fn
+
+    def on_shutdown[Call_T: AsyncCallable[[], None]](self, fn: Call_T) -> Call_T:
+        self.shutdowns.append(fn)
+        return fn
+
+    def on_context[Ctx_T: Callable[[], AsyncGenerator[Any, None]]](self, fn: Ctx_T) -> Ctx_T:
+        self.contexts.append(fn)
+        return fn