mirror of
https://github.com/zhenxun-org/zhenxun_bot.git
synced 2025-12-15 14:22:55 +08:00
58 lines
1.7 KiB
Python
58 lines
1.7 KiB
Python
from collections.abc import Callable
|
|
from typing import ClassVar
|
|
|
|
import nonebot
|
|
from nonebot.utils import is_coroutine_callable
|
|
|
|
from zhenxun.services.log import logger
|
|
from zhenxun.utils.enum import PriorityLifecycle
|
|
from zhenxun.utils.exception import HookPriorityException
|
|
|
|
driver = nonebot.get_driver()
|
|
|
|
|
|
class PriorityLifecycle:
|
|
_data: ClassVar[dict[PriorityLifecycle, dict[int, list[Callable]]]] = {}
|
|
|
|
@classmethod
|
|
def add(cls, hook_type: PriorityLifecycle, func: Callable, priority: int = 5):
|
|
if hook_type not in cls._data:
|
|
cls._data[hook_type] = {}
|
|
if priority not in cls._data[hook_type]:
|
|
cls._data[hook_type][priority] = []
|
|
cls._data[hook_type][priority].append(func)
|
|
|
|
@classmethod
|
|
def on_startup(cls, *, priority: int = 5):
|
|
def wrapper(func):
|
|
cls.add(PriorityLifecycle.STARTUP, func, priority)
|
|
return func
|
|
|
|
return wrapper
|
|
|
|
@classmethod
|
|
def on_shutdown(cls, *, priority: int = 5):
|
|
def wrapper(func):
|
|
cls.add(PriorityLifecycle.SHUTDOWN, func, priority)
|
|
return func
|
|
|
|
return wrapper
|
|
|
|
|
|
@driver.on_startup
|
|
async def _():
|
|
priority_data = PriorityLifecycle._data.get(PriorityLifecycle.STARTUP)
|
|
if not priority_data:
|
|
return
|
|
priority_list = sorted(priority_data.keys())
|
|
priority = 0
|
|
try:
|
|
for priority in priority_list:
|
|
for func in priority_data[priority]:
|
|
if is_coroutine_callable(func):
|
|
await func()
|
|
else:
|
|
func()
|
|
except HookPriorityException as e:
|
|
logger.error(f"打断优先级 [{priority}] on_startup 方法. {type(e)}: {e}")
|