mirror of
https://github.com/nonebot/nonebot2.git
synced 2025-09-07 12:36:47 +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:
@ -1,46 +0,0 @@
|
||||
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._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
|
||||
|
||||
@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)
|
||||
|
||||
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()
|
@ -34,8 +34,6 @@ from nonebot.drivers import Request as BaseRequest
|
||||
from nonebot.drivers import WebSocket as BaseWebSocket
|
||||
from nonebot.drivers import HTTPServerSetup, WebSocketServerSetup
|
||||
|
||||
from ._lifespan import LIFESPAN_FUNC, Lifespan
|
||||
|
||||
try:
|
||||
import uvicorn
|
||||
from fastapi.responses import Response
|
||||
@ -97,8 +95,6 @@ class Driver(BaseDriver, ASGIMixin):
|
||||
|
||||
self.fastapi_config: Config = Config(**config.dict())
|
||||
|
||||
self._lifespan = Lifespan()
|
||||
|
||||
self._server_app = FastAPI(
|
||||
lifespan=self._lifespan_manager,
|
||||
openapi_url=self.fastapi_config.fastapi_openapi_url,
|
||||
@ -155,14 +151,6 @@ class Driver(BaseDriver, ASGIMixin):
|
||||
name=setup.name,
|
||||
)
|
||||
|
||||
@override
|
||||
def on_startup(self, func: LIFESPAN_FUNC) -> LIFESPAN_FUNC:
|
||||
return self._lifespan.on_startup(func)
|
||||
|
||||
@override
|
||||
def on_shutdown(self, func: LIFESPAN_FUNC) -> LIFESPAN_FUNC:
|
||||
return self._lifespan.on_shutdown(func)
|
||||
|
||||
@contextlib.asynccontextmanager
|
||||
async def _lifespan_manager(self, app: FastAPI):
|
||||
await self._lifespan.startup()
|
||||
|
@ -19,8 +19,6 @@ from nonebot.consts import WINDOWS
|
||||
from nonebot.config import Env, Config
|
||||
from nonebot.drivers import Driver as BaseDriver
|
||||
|
||||
from ._lifespan import LIFESPAN_FUNC, Lifespan
|
||||
|
||||
HANDLED_SIGNALS = (
|
||||
signal.SIGINT, # Unix signal 2. Sent by Ctrl+C.
|
||||
signal.SIGTERM, # Unix signal 15. Sent by `kill <pid>`.
|
||||
@ -35,8 +33,6 @@ class Driver(BaseDriver):
|
||||
def __init__(self, env: Env, config: Config):
|
||||
super().__init__(env, config)
|
||||
|
||||
self._lifespan = Lifespan()
|
||||
|
||||
self.should_exit: asyncio.Event = asyncio.Event()
|
||||
self.force_exit: bool = False
|
||||
|
||||
@ -52,16 +48,6 @@ class Driver(BaseDriver):
|
||||
"""none driver 使用的 logger"""
|
||||
return logger
|
||||
|
||||
@override
|
||||
def on_startup(self, func: LIFESPAN_FUNC) -> LIFESPAN_FUNC:
|
||||
"""注册一个启动时执行的函数"""
|
||||
return self._lifespan.on_startup(func)
|
||||
|
||||
@override
|
||||
def on_shutdown(self, func: LIFESPAN_FUNC) -> LIFESPAN_FUNC:
|
||||
"""注册一个停止时执行的函数"""
|
||||
return self._lifespan.on_shutdown(func)
|
||||
|
||||
@override
|
||||
def run(self, *args, **kwargs):
|
||||
"""启动 none driver"""
|
||||
|
@ -18,18 +18,7 @@ FrontMatter:
|
||||
import asyncio
|
||||
from functools import wraps
|
||||
from typing_extensions import override
|
||||
from typing import (
|
||||
Any,
|
||||
Dict,
|
||||
List,
|
||||
Tuple,
|
||||
Union,
|
||||
TypeVar,
|
||||
Callable,
|
||||
Optional,
|
||||
Coroutine,
|
||||
cast,
|
||||
)
|
||||
from typing import Any, Dict, List, Tuple, Union, Optional, cast
|
||||
|
||||
from pydantic import BaseSettings
|
||||
|
||||
@ -57,8 +46,6 @@ except ModuleNotFoundError as e: # pragma: no cover
|
||||
"Install with pip: `pip install nonebot2[quart]`"
|
||||
) from e
|
||||
|
||||
_AsyncCallable = TypeVar("_AsyncCallable", bound=Callable[..., Coroutine])
|
||||
|
||||
|
||||
def catch_closed(func):
|
||||
@wraps(func)
|
||||
@ -102,6 +89,8 @@ class Driver(BaseDriver, ASGIMixin):
|
||||
self._server_app = Quart(
|
||||
self.__class__.__qualname__, **self.quart_config.quart_extra
|
||||
)
|
||||
self._server_app.before_serving(self._lifespan.startup)
|
||||
self._server_app.after_serving(self._lifespan.shutdown)
|
||||
|
||||
@property
|
||||
@override
|
||||
@ -150,16 +139,6 @@ class Driver(BaseDriver, ASGIMixin):
|
||||
view_func=_handle,
|
||||
)
|
||||
|
||||
@override
|
||||
def on_startup(self, func: _AsyncCallable) -> _AsyncCallable:
|
||||
"""参考文档: [`Startup and Shutdown`](https://pgjones.gitlab.io/quart/how_to_guides/startup_shutdown.html)"""
|
||||
return self.server_app.before_serving(func) # type: ignore
|
||||
|
||||
@override
|
||||
def on_shutdown(self, func: _AsyncCallable) -> _AsyncCallable:
|
||||
"""参考文档: [`Startup and Shutdown`](https://pgjones.gitlab.io/quart/how_to_guides/startup_shutdown.html)"""
|
||||
return self.server_app.after_serving(func) # type: ignore
|
||||
|
||||
@override
|
||||
def run(
|
||||
self,
|
||||
|
Reference in New Issue
Block a user