🐛 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>

nonebot/internal/driver/_lifespan.py

@@ -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()

nonebot/internal/driver/abstract.py

@@ -2,7 +2,7 @@ import abc
 import asyncio
 from typing_extensions import TypeAlias
 from contextlib import AsyncExitStack, asynccontextmanager
-from typing import TYPE_CHECKING, Any, Set, Dict, Type, Callable, AsyncGenerator
+from typing import TYPE_CHECKING, Any, Set, Dict, Type, AsyncGenerator
 
 from nonebot.log import logger
 from nonebot.config import Env, Config
@@ -16,6 +16,7 @@ from nonebot.typing import (
     T_BotDisconnectionHook,
 )
 
+from ._lifespan import LIFESPAN_FUNC, Lifespan
 from .model import Request, Response, WebSocket, HTTPServerSetup, WebSocketServerSetup
 
 if TYPE_CHECKING:
@@ -49,6 +50,7 @@ class Driver(abc.ABC):
         """全局配置对象"""
         self._bots: Dict[str, "Bot"] = {}
         self._bot_tasks: Set[asyncio.Task] = set()
+        self._lifespan = Lifespan()
 
     def __repr__(self) -> str:
         return (
@@ -100,15 +102,13 @@ class Driver(abc.ABC):
 
         self.on_shutdown(self._cleanup)
 
-    @abc.abstractmethod
-    def on_startup(self, func: Callable) -> Callable:
-        """注册一个在驱动器启动时执行的函数"""
-        raise NotImplementedError
+    def on_startup(self, func: LIFESPAN_FUNC) -> LIFESPAN_FUNC:
+        """注册一个启动时执行的函数"""
+        return self._lifespan.on_startup(func)
 
-    @abc.abstractmethod
-    def on_shutdown(self, func: Callable) -> Callable:
-        """注册一个在驱动器停止时执行的函数"""
-        raise NotImplementedError
+    def on_shutdown(self, func: LIFESPAN_FUNC) -> LIFESPAN_FUNC:
+        """注册一个停止时执行的函数"""
+        return self._lifespan.on_shutdown(func)
 
     @classmethod
     def on_bot_connect(cls, func: T_BotConnectionHook) -> T_BotConnectionHook: