2024-05-13 23:06:08 +08:00

58 lines
2.6 KiB
Python

from nonebot.log import logger
from ..config import config
from .scheduler import Scheduler
from ..utils import SchedulerConfig
from ..config.db_model import Target
from ..types import Target as T_Target
from ..platform import platform_manager
from ..plugin_config import plugin_config
scheduler_dict: dict[type[SchedulerConfig], Scheduler] = {}
async def init_scheduler():
_schedule_class_dict: dict[type[SchedulerConfig], list[Target]] = {}
_schedule_class_platform_dict: dict[type[SchedulerConfig], list[str]] = {}
for platform in platform_manager.values():
scheduler_config = platform.scheduler
if not hasattr(scheduler_config, "name") or not scheduler_config.name:
scheduler_config.name = f"AnonymousScheduleConfig[{platform.platform_name}]"
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] = list(targets)
else:
_schedule_class_dict[scheduler_config].extend(targets)
if scheduler_config not in _schedule_class_platform_dict:
_schedule_class_platform_dict[scheduler_config] = [platform_name]
else:
_schedule_class_platform_dict[scheduler_config].append(platform_name)
for scheduler_config, target_list in _schedule_class_dict.items():
if not plugin_config.bison_use_browser and scheduler_config.require_browser:
logger.warning(f"{scheduler_config.name} requires browser, it will not schedule.")
continue
schedulable_args = []
for target in target_list:
schedulable_args.append(
(target.platform_name, T_Target(target.target), platform_manager[target.platform_name].use_batch)
)
platform_name_list = _schedule_class_platform_dict[scheduler_config]
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):
platform = platform_manager[platform_name]
scheduler_obj = scheduler_dict[platform.scheduler]
scheduler_obj.insert_new_schedulable(platform_name, target)
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)