nonebot2/nonebot/drivers/none.py
Bryan不可思议 8f3f385cb6
🐛 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>
2023-12-10 18:12:10 +08:00

142 lines
4.0 KiB
Python

"""None 驱动适配
:::tip 提示
本驱动不支持任何服务器或客户端连接
:::
FrontMatter:
sidebar_position: 6
description: nonebot.drivers.none 模块
"""
import signal
import asyncio
import threading
from typing_extensions import override
from nonebot.log import logger
from nonebot.consts import WINDOWS
from nonebot.config import Env, Config
from nonebot.drivers import Driver as BaseDriver
HANDLED_SIGNALS = (
signal.SIGINT, # Unix signal 2. Sent by Ctrl+C.
signal.SIGTERM, # Unix signal 15. Sent by `kill <pid>`.
)
if WINDOWS: # pragma: py-win32
HANDLED_SIGNALS += (signal.SIGBREAK,) # Windows signal 21. Sent by Ctrl+Break.
class Driver(BaseDriver):
"""None 驱动框架"""
def __init__(self, env: Env, config: Config):
super().__init__(env, config)
self.should_exit: asyncio.Event = asyncio.Event()
self.force_exit: bool = False
@property
@override
def type(self) -> str:
"""驱动名称: `none`"""
return "none"
@property
@override
def logger(self):
"""none driver 使用的 logger"""
return logger
@override
def run(self, *args, **kwargs):
"""启动 none driver"""
super().run(*args, **kwargs)
loop = asyncio.get_event_loop()
loop.run_until_complete(self._serve())
async def _serve(self):
self._install_signal_handlers()
await self._startup()
if self.should_exit.is_set():
return
await self._main_loop()
await self._shutdown()
async def _startup(self):
try:
await self._lifespan.startup()
except Exception as e:
logger.opt(colors=True, exception=e).error(
"<r><bg #f8bbd0>Error when running startup function. "
"Ignored!</bg #f8bbd0></r>"
)
logger.info("Application startup completed.")
async def _main_loop(self):
await self.should_exit.wait()
async def _shutdown(self):
logger.info("Shutting down")
logger.info("Waiting for application shutdown.")
try:
await self._lifespan.shutdown()
except Exception as e:
logger.opt(colors=True, exception=e).error(
"<r><bg #f8bbd0>Error when running shutdown function. "
"Ignored!</bg #f8bbd0></r>"
)
for task in asyncio.all_tasks():
if task is not asyncio.current_task() and not task.done():
task.cancel()
await asyncio.sleep(0.1)
tasks = [t for t in asyncio.all_tasks() if t is not asyncio.current_task()]
if tasks and not self.force_exit:
logger.info("Waiting for tasks to finish. (CTRL+C to force quit)")
while tasks and not self.force_exit:
await asyncio.sleep(0.1)
tasks = [t for t in asyncio.all_tasks() if t is not asyncio.current_task()]
for task in tasks:
task.cancel()
await asyncio.gather(*tasks, return_exceptions=True)
logger.info("Application shutdown complete.")
loop = asyncio.get_event_loop()
loop.stop()
def _install_signal_handlers(self) -> None:
if threading.current_thread() is not threading.main_thread():
# Signals can only be listened to from the main thread.
return
loop = asyncio.get_event_loop()
try:
for sig in HANDLED_SIGNALS:
loop.add_signal_handler(sig, self._handle_exit, sig, None)
except NotImplementedError:
# Windows
for sig in HANDLED_SIGNALS:
signal.signal(sig, self._handle_exit)
def _handle_exit(self, sig, frame):
self.exit(force=self.should_exit.is_set())
def exit(self, force: bool = False):
"""退出 none driver
参数:
force: 强制退出
"""
if not self.should_exit.is_set():
self.should_exit.set()
if force:
self.force_exit = True