60 lines
1.9 KiB
Python
60 lines
1.9 KiB
Python
from collections.abc import AsyncGenerator, Callable
|
|
from typing import Any, TypeGuard, override
|
|
|
|
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:
|
|
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()
|
|
|
|
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
|