mirror of
https://github.com/nonebot/nonebot2.git
synced 2025-07-16 19:11:00 +00:00
🐛 Fix: 新增 Lifespan._on_ready()
供适配器使用 (#2483)
Co-authored-by: Ju4tCode <42488585+yanyongyu@users.noreply.github.com> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
This commit is contained in:
54
nonebot/internal/driver/_lifespan.py
Normal file
54
nonebot/internal/driver/_lifespan.py
Normal file
@ -0,0 +1,54 @@
|
||||
from typing_extensions import TypeAlias
|
||||
from typing import Any, List, Union, Callable, Awaitable, cast
|
||||
|
||||
from nonebot.utils import run_sync, is_coroutine_callable
|
||||
|
||||
SYNC_LIFESPAN_FUNC: TypeAlias = Callable[[], Any]
|
||||
ASYNC_LIFESPAN_FUNC: TypeAlias = Callable[[], Awaitable[Any]]
|
||||
LIFESPAN_FUNC: TypeAlias = Union[SYNC_LIFESPAN_FUNC, ASYNC_LIFESPAN_FUNC]
|
||||
|
||||
|
||||
class Lifespan:
|
||||
def __init__(self) -> None:
|
||||
self._startup_funcs: List[LIFESPAN_FUNC] = []
|
||||
self._ready_funcs: List[LIFESPAN_FUNC] = []
|
||||
self._shutdown_funcs: List[LIFESPAN_FUNC] = []
|
||||
|
||||
def on_startup(self, func: LIFESPAN_FUNC) -> LIFESPAN_FUNC:
|
||||
self._startup_funcs.append(func)
|
||||
return func
|
||||
|
||||
def on_shutdown(self, func: LIFESPAN_FUNC) -> LIFESPAN_FUNC:
|
||||
self._shutdown_funcs.append(func)
|
||||
return func
|
||||
|
||||
def on_ready(self, func: LIFESPAN_FUNC) -> LIFESPAN_FUNC:
|
||||
self._ready_funcs.append(func)
|
||||
return func
|
||||
|
||||
@staticmethod
|
||||
async def _run_lifespan_func(
|
||||
funcs: List[LIFESPAN_FUNC],
|
||||
) -> None:
|
||||
for func in funcs:
|
||||
if is_coroutine_callable(func):
|
||||
await cast(ASYNC_LIFESPAN_FUNC, func)()
|
||||
else:
|
||||
await run_sync(cast(SYNC_LIFESPAN_FUNC, func))()
|
||||
|
||||
async def startup(self) -> None:
|
||||
if self._startup_funcs:
|
||||
await self._run_lifespan_func(self._startup_funcs)
|
||||
|
||||
if self._ready_funcs:
|
||||
await self._run_lifespan_func(self._ready_funcs)
|
||||
|
||||
async def shutdown(self) -> None:
|
||||
if self._shutdown_funcs:
|
||||
await self._run_lifespan_func(self._shutdown_funcs)
|
||||
|
||||
async def __aenter__(self) -> None:
|
||||
await self.startup()
|
||||
|
||||
async def __aexit__(self, exc_type, exc_val, exc_tb) -> None:
|
||||
await self.shutdown()
|
Reference in New Issue
Block a user