diff --git a/bot.py b/bot.py index 52cd29fc..aa047a71 100644 --- a/bot.py +++ b/bot.py @@ -14,9 +14,9 @@ driver.register_adapter(OneBotV11Adapter) # driver.register_adapter(DoDoAdapter) # driver.register_adapter(DiscordAdapter) -from zhenxun.services.db_context import disconnect, init +from zhenxun.services.db_context import disconnect -driver.on_startup(init) +# driver.on_startup(init) driver.on_shutdown(disconnect) # nonebot.load_builtin_plugins("echo") diff --git a/zhenxun/builtin_plugins/__init__.py b/zhenxun/builtin_plugins/__init__.py index fbaeb280..7bc82012 100644 --- a/zhenxun/builtin_plugins/__init__.py +++ b/zhenxun/builtin_plugins/__init__.py @@ -16,6 +16,7 @@ from zhenxun.models.sign_user import SignUser from zhenxun.models.user_console import UserConsole from zhenxun.services.log import logger from zhenxun.utils.decorator.shop import shop_register +from zhenxun.utils.manager.priority_manager import HookPriorityManager from zhenxun.utils.manager.resource_manager import ResourceManager from zhenxun.utils.platform import PlatformUtils @@ -70,7 +71,7 @@ from public.bag_users t1 """ -@driver.on_startup +@HookPriorityManager.on_startup() async def _(): await ResourceManager.init_resources() """签到与用户的数据迁移""" diff --git a/zhenxun/builtin_plugins/admin/welcome_message/data_source.py b/zhenxun/builtin_plugins/admin/welcome_message/data_source.py index 811a8dda..3b56a981 100644 --- a/zhenxun/builtin_plugins/admin/welcome_message/data_source.py +++ b/zhenxun/builtin_plugins/admin/welcome_message/data_source.py @@ -14,6 +14,7 @@ from zhenxun.services.log import logger from zhenxun.utils._build_image import BuildImage from zhenxun.utils._image_template import ImageTemplate from zhenxun.utils.http_utils import AsyncHttpx +from zhenxun.utils.manager.priority_manager import HookPriorityManager from zhenxun.utils.platform import PlatformUtils BASE_PATH = DATA_PATH / "welcome_message" @@ -87,7 +88,7 @@ def migrate(path: Path): json.dump(new_data, f, ensure_ascii=False, indent=4) -@driver.on_startup +@HookPriorityManager.on_startup() def _(): """数据迁移 diff --git a/zhenxun/builtin_plugins/init/init_config.py b/zhenxun/builtin_plugins/init/init_config.py index 112d29de..b81df4fd 100644 --- a/zhenxun/builtin_plugins/init/init_config.py +++ b/zhenxun/builtin_plugins/init/init_config.py @@ -11,6 +11,7 @@ from zhenxun.configs.config import Config from zhenxun.configs.path_config import DATA_PATH from zhenxun.configs.utils import RegisterConfig from zhenxun.services.log import logger +from zhenxun.utils.manager.priority_manager import HookPriorityManager _yaml = YAML(pure=True) _yaml.allow_unicode = True @@ -102,7 +103,7 @@ def _generate_simple_config(exists_module: list[str]): temp_file.unlink() -@driver.on_startup +@HookPriorityManager.on_startup() def _(): """ 初始化插件数据配置 diff --git a/zhenxun/builtin_plugins/init/init_plugin.py b/zhenxun/builtin_plugins/init/init_plugin.py index dbeddb54..510fc5a2 100644 --- a/zhenxun/builtin_plugins/init/init_plugin.py +++ b/zhenxun/builtin_plugins/init/init_plugin.py @@ -20,6 +20,7 @@ from zhenxun.utils.enum import ( PluginLimitType, PluginType, ) +from zhenxun.utils.manager.priority_manager import HookPriorityManager from .manager import manager @@ -95,7 +96,7 @@ async def _handle_setting( ) -@driver.on_startup +@HookPriorityManager.on_startup() async def _(): """ 初始化插件数据配置 diff --git a/zhenxun/builtin_plugins/init/init_task.py b/zhenxun/builtin_plugins/init/init_task.py index cead7d72..50475ccd 100644 --- a/zhenxun/builtin_plugins/init/init_task.py +++ b/zhenxun/builtin_plugins/init/init_task.py @@ -10,6 +10,7 @@ from zhenxun.models.group_console import GroupConsole from zhenxun.models.task_info import TaskInfo from zhenxun.services.log import logger from zhenxun.utils.common_utils import CommonUtils +from zhenxun.utils.manager.priority_manager import HookPriorityManager driver: Driver = nonebot.get_driver() @@ -132,7 +133,7 @@ async def create_schedule(task: Task): logger.error(f"动态创建定时任务 {task.name}({task.module}) 失败", e=e) -@driver.on_startup +@HookPriorityManager.on_startup() async def _(): """ 初始化插件数据配置 diff --git a/zhenxun/builtin_plugins/scripts.py b/zhenxun/builtin_plugins/scripts.py index 27705301..02516b64 100644 --- a/zhenxun/builtin_plugins/scripts.py +++ b/zhenxun/builtin_plugins/scripts.py @@ -10,11 +10,12 @@ from zhenxun.configs.path_config import TEXT_PATH from zhenxun.models.group_console import GroupConsole from zhenxun.services.log import logger from zhenxun.utils.http_utils import AsyncHttpx +from zhenxun.utils.manager.priority_manager import HookPriorityManager driver: Driver = nonebot.get_driver() -@driver.on_startup +@HookPriorityManager.on_startup() async def update_city(): """ 部分插件需要中国省份城市 @@ -61,7 +62,7 @@ async def _(): await update_city() -@driver.on_startup +@HookPriorityManager.on_startup() async def _(): """开启/禁用插件格式修改""" _, is_create = await GroupConsole.get_or_create(group_id=133133133) diff --git a/zhenxun/builtin_plugins/sign_in/__init__.py b/zhenxun/builtin_plugins/sign_in/__init__.py index 0b48a0e7..0986e476 100644 --- a/zhenxun/builtin_plugins/sign_in/__init__.py +++ b/zhenxun/builtin_plugins/sign_in/__init__.py @@ -10,7 +10,6 @@ from nonebot_plugin_alconna import ( store_true, ) from nonebot_plugin_apscheduler import scheduler -from nonebot_plugin_uninfo import Uninfo from zhenxun.configs.utils import ( Command, @@ -23,7 +22,7 @@ from zhenxun.utils.depends import UserName from zhenxun.utils.message import MessageUtils from ._data_source import SignManage -from .goods_register import driver # noqa: F401 +from .goods_register import Uninfo from .utils import clear_sign_data_pic __plugin_meta__ = PluginMetadata( diff --git a/zhenxun/builtin_plugins/sign_in/goods_register.py b/zhenxun/builtin_plugins/sign_in/goods_register.py index f7a65359..6c8e39bb 100644 --- a/zhenxun/builtin_plugins/sign_in/goods_register.py +++ b/zhenxun/builtin_plugins/sign_in/goods_register.py @@ -1,7 +1,6 @@ from decimal import Decimal import nonebot -from nonebot.drivers import Driver from nonebot_plugin_uninfo import Uninfo from zhenxun.models.sign_user import SignUser @@ -9,14 +8,7 @@ from zhenxun.models.user_console import UserConsole from zhenxun.utils.decorator.shop import shop_register from zhenxun.utils.platform import PlatformUtils -driver: Driver = nonebot.get_driver() - - -# @driver.on_startup -# async def _(): -# """ -# 导入内置的三个商品 -# """ +driver = nonebot.get_driver() @shop_register( diff --git a/zhenxun/builtin_plugins/sign_in/utils.py b/zhenxun/builtin_plugins/sign_in/utils.py index 9faf1120..84c0d8cd 100644 --- a/zhenxun/builtin_plugins/sign_in/utils.py +++ b/zhenxun/builtin_plugins/sign_in/utils.py @@ -16,6 +16,7 @@ from zhenxun.models.sign_log import SignLog from zhenxun.models.sign_user import SignUser from zhenxun.utils.http_utils import AsyncHttpx from zhenxun.utils.image_utils import BuildImage +from zhenxun.utils.manager.priority_manager import HookPriorityManager from zhenxun.utils.platform import PlatformUtils from .config import ( @@ -54,7 +55,7 @@ LG_MESSAGE = [ ] -@driver.on_startup +@HookPriorityManager.on_startup() async def init_image(): SIGN_RESOURCE_PATH.mkdir(parents=True, exist_ok=True) SIGN_TODAY_CARD_PATH.mkdir(exist_ok=True, parents=True) diff --git a/zhenxun/builtin_plugins/web_ui/__init__.py b/zhenxun/builtin_plugins/web_ui/__init__.py index d8d71025..5bba583d 100644 --- a/zhenxun/builtin_plugins/web_ui/__init__.py +++ b/zhenxun/builtin_plugins/web_ui/__init__.py @@ -10,6 +10,7 @@ from zhenxun.configs.config import Config as gConfig from zhenxun.configs.utils import PluginExtraData, RegisterConfig from zhenxun.services.log import logger, logger_ from zhenxun.utils.enum import PluginType +from zhenxun.utils.manager.priority_manager import HookPriorityManager from .api.logs import router as ws_log_routes from .api.logs.log_manager import LOG_STORAGE @@ -91,7 +92,7 @@ WsApiRouter.include_router(status_routes) WsApiRouter.include_router(chat_routes) -@driver.on_startup +@HookPriorityManager.on_startup() async def _(): try: diff --git a/zhenxun/builtin_plugins/web_ui/api/tabs/dashboard/data_source.py b/zhenxun/builtin_plugins/web_ui/api/tabs/dashboard/data_source.py index 6c312db3..535cac22 100644 --- a/zhenxun/builtin_plugins/web_ui/api/tabs/dashboard/data_source.py +++ b/zhenxun/builtin_plugins/web_ui/api/tabs/dashboard/data_source.py @@ -13,6 +13,7 @@ from zhenxun.models.bot_connect_log import BotConnectLog from zhenxun.models.chat_history import ChatHistory from zhenxun.models.statistics import Statistics from zhenxun.services.log import logger +from zhenxun.utils.manager.priority_manager import HookPriorityManager from zhenxun.utils.platform import PlatformUtils from ....base_model import BaseResultModel, QueryModel @@ -31,7 +32,7 @@ driver: Driver = nonebot.get_driver() CONNECT_TIME = 0 -@driver.on_startup +@HookPriorityManager.on_startup() async def _(): global CONNECT_TIME CONNECT_TIME = int(time.time()) diff --git a/zhenxun/builtin_plugins/web_ui/api/tabs/database/__init__.py b/zhenxun/builtin_plugins/web_ui/api/tabs/database/__init__.py index ed6df2e3..32ba0df8 100644 --- a/zhenxun/builtin_plugins/web_ui/api/tabs/database/__init__.py +++ b/zhenxun/builtin_plugins/web_ui/api/tabs/database/__init__.py @@ -8,6 +8,7 @@ from zhenxun.configs.config import BotConfig from zhenxun.models.plugin_info import PluginInfo from zhenxun.models.task_info import TaskInfo from zhenxun.services.log import logger +from zhenxun.utils.manager.priority_manager import HookPriorityManager from ....base_model import BaseResultModel, QueryModel, Result from ....utils import authentication @@ -21,7 +22,7 @@ router = APIRouter(prefix="/database") driver: Driver = nonebot.get_driver() -@driver.on_startup +@HookPriorityManager.on_startup() async def _(): for plugin in nonebot.get_loaded_plugins(): module = plugin.name diff --git a/zhenxun/services/db_context.py b/zhenxun/services/db_context.py index 9a44fa74..76c4cf98 100644 --- a/zhenxun/services/db_context.py +++ b/zhenxun/services/db_context.py @@ -1,9 +1,12 @@ +import nonebot from nonebot.utils import is_coroutine_callable from tortoise import Tortoise from tortoise.connection import connections from tortoise.models import Model as Model_ from zhenxun.configs.config import BotConfig +from zhenxun.utils.exception import HookPriorityException +from zhenxun.utils.manager.priority_manager import HookPriorityManager from .log import logger @@ -11,6 +14,9 @@ SCRIPT_METHOD = [] MODELS: list[str] = [] +driver = nonebot.get_driver() + + class Model(Model_): """ 自动添加模块 @@ -26,7 +32,7 @@ class Model(Model_): SCRIPT_METHOD.append((cls.__module__, func)) -class DbUrlIsNode(Exception): +class DbUrlIsNode(HookPriorityException): """ 数据库链接地址为空 """ @@ -42,9 +48,18 @@ class DbConnectError(Exception): pass +@HookPriorityManager.on_startup(1) async def init(): if not BotConfig.db_url: - raise DbUrlIsNode("数据库配置为空,请在.env.dev中配置DB_URL...") + # raise DbUrlIsNode("数据库配置为空,请在.env.dev中配置DB_URL...") + error = f""" +*********************************************************************** +**********************************配置为空****************************** +请打开WebUi进行基础配置,配置地址:http://{driver.config.host}:{driver.config.port} +*********************************************************************** +*********************************************************************** + """ + raise DbUrlIsNode("\n" + error.strip()) try: await Tortoise.init( db_url=BotConfig.db_url, diff --git a/zhenxun/services/plugin_init.py b/zhenxun/services/plugin_init.py index 159e042c..4993950f 100644 --- a/zhenxun/services/plugin_init.py +++ b/zhenxun/services/plugin_init.py @@ -6,6 +6,7 @@ from nonebot.utils import is_coroutine_callable from pydantic import BaseModel from zhenxun.services.log import logger +from zhenxun.utils.manager.priority_manager import HookPriorityManager driver = nonebot.get_driver() @@ -100,6 +101,6 @@ class PluginInitManager: logger.error(f"执行: {module_path}:remove 失败", e=e) -@driver.on_startup +@HookPriorityManager.on_startup() async def _(): await PluginInitManager.install_all() diff --git a/zhenxun/utils/enum.py b/zhenxun/utils/enum.py index c0b79342..a68d8918 100644 --- a/zhenxun/utils/enum.py +++ b/zhenxun/utils/enum.py @@ -1,6 +1,13 @@ from strenum import StrEnum +class HookPriorityType(StrEnum): + STARTUP = "STARTUP" + """启动""" + SHUTDOWN = "SHUTDOWN" + """关闭""" + + class GoldHandle(StrEnum): """ 金币处理 diff --git a/zhenxun/utils/exception.py b/zhenxun/utils/exception.py index db8c0656..8ec925ec 100644 --- a/zhenxun/utils/exception.py +++ b/zhenxun/utils/exception.py @@ -1,3 +1,15 @@ +class HookPriorityException(BaseException): + """ + 钩子优先级异常 + """ + + def __init__(self, info: str = "") -> None: + self.info = info + + def __str__(self) -> str: + return self.info + + class NotFoundError(Exception): """ 未发现 diff --git a/zhenxun/utils/manager/message_manager.py b/zhenxun/utils/manager/message_manager.py index e714c8d8..ee34369d 100644 --- a/zhenxun/utils/manager/message_manager.py +++ b/zhenxun/utils/manager/message_manager.py @@ -22,6 +22,4 @@ class MessageManager: @classmethod def get(cls, uid: str) -> list[str]: - if uid in cls.data: - return cls.data[uid] - return [] + return cls.data[uid] if uid in cls.data else [] diff --git a/zhenxun/utils/manager/priority_manager.py b/zhenxun/utils/manager/priority_manager.py new file mode 100644 index 00000000..0cb454b3 --- /dev/null +++ b/zhenxun/utils/manager/priority_manager.py @@ -0,0 +1,57 @@ +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 HookPriorityType +from zhenxun.utils.exception import HookPriorityException + +driver = nonebot.get_driver() + + +class HookPriorityManager: + _data: ClassVar[dict[HookPriorityType, dict[int, list[Callable]]]] = {} + + @classmethod + def add(cls, hook_type: HookPriorityType, 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(HookPriorityType.STARTUP, func, priority) + return func + + return wrapper + + @classmethod + def on_shutdown(cls, priority: int = 5): + def wrapper(func): + cls.add(HookPriorityType.SHUTDOWN, func, priority) + return func + + return wrapper + + +@driver.on_startup +async def _(): + priority_data = HookPriorityManager._data.get(HookPriorityType.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}")