From 010e6335f27ed36fc5b8df9c4509ac83a9a1ddc6 Mon Sep 17 00:00:00 2001 From: felinae98 <731499577@qq.com> Date: Sun, 12 Mar 2023 17:43:32 +0800 Subject: [PATCH] :bug: register db hook after init scheduler --- nonebot_bison/config/db_config.py | 23 ++++++++++++++--------- nonebot_bison/scheduler/manager.py | 8 +++----- 2 files changed, 17 insertions(+), 14 deletions(-) diff --git a/nonebot_bison/config/db_config.py b/nonebot_bison/config/db_config.py index 7f43416..59922c6 100644 --- a/nonebot_bison/config/db_config.py +++ b/nonebot_bison/config/db_config.py @@ -1,3 +1,4 @@ +import asyncio from collections import defaultdict from datetime import datetime, time from typing import Awaitable, Callable, Optional, Sequence @@ -28,14 +29,14 @@ class SubscribeDupException(Exception): class DBConfig: def __init__(self): - self.add_target_hook: Optional[Callable[[str, T_Target], Awaitable]] = None - self.delete_target_hook: Optional[Callable[[str, T_Target], Awaitable]] = None + self.add_target_hook: list[Callable[[str, T_Target], Awaitable]] = [] + self.delete_target_hook: list[Callable[[str, T_Target], Awaitable]] = [] def register_add_target_hook(self, fun: Callable[[str, T_Target], Awaitable]): - self.add_target_hook = fun + self.add_target_hook.append(fun) def register_delete_target_hook(self, fun: Callable[[str, T_Target], Awaitable]): - self.delete_target_hook = fun + self.delete_target_hook.append(fun) async def add_subscribe( self, @@ -65,8 +66,9 @@ class DBConfig: db_target = Target( target=target, platform_name=platform_name, target_name=target_name ) - if self.add_target_hook: - await self.add_target_hook(platform_name, target) + await asyncio.gather( + *[hook(platform_name, target) for hook in self.add_target_hook] + ) else: db_target.target_name = target_name subscribe = Subscribe( @@ -118,9 +120,12 @@ class DBConfig: ) if target_count == 0: # delete empty target - # await session.delete(target_obj) - if self.delete_target_hook: - await self.delete_target_hook(platform_name, T_Target(target)) + await asyncio.gather( + *[ + hook(platform_name, T_Target(target)) + for hook in self.delete_target_hook + ] + ) await session.commit() async def update_subscribe( diff --git a/nonebot_bison/scheduler/manager.py b/nonebot_bison/scheduler/manager.py index f41ac51..9c48676 100644 --- a/nonebot_bison/scheduler/manager.py +++ b/nonebot_bison/scheduler/manager.py @@ -21,7 +21,7 @@ async def init_scheduler(): platform_name = platform.platform_name targets = await config.get_platform_target(platform_name) if scheduler_config not in _schedule_class_dict: - _schedule_class_dict[scheduler_config] = targets + _schedule_class_dict[scheduler_config] = list(targets) else: _schedule_class_dict[scheduler_config].extend(targets) if scheduler_config not in _schedule_class_platform_dict: @@ -36,6 +36,8 @@ async def init_scheduler(): scheduler_dict[scheduler_config] = Scheduler( scheduler_config, schedulable_args, platform_name_list ) + config.register_add_target_hook(handle_insert_new_target) + config.register_delete_target_hook(handle_delete_target) async def handle_insert_new_target(platform_name: str, target: T_Target): @@ -48,7 +50,3 @@ async def handle_delete_target(platform_name: str, target: T_Target): platform = platform_manager[platform_name] scheduler_obj = scheduler_dict[platform.scheduler] scheduler_obj.delete_schedulable(platform_name, target) - - -config.register_add_target_hook(handle_insert_new_target) -config.register_delete_target_hook(handle_delete_target)