From b613abeb8bc6975e2fe59463ff4541c503eeafba Mon Sep 17 00:00:00 2001 From: LWR Date: Sat, 29 Oct 2022 17:46:32 +0800 Subject: [PATCH] feat: Dict datasource support, Add StarBot model --- starbot/core/bot.py | 58 ++++++++++++ starbot/core/datasource.py | 114 +++++++++++++++++++++++ starbot/exception/DataSourceException.py | 15 +++ starbot/exception/__init__.py | 1 + 4 files changed, 188 insertions(+) create mode 100644 starbot/core/bot.py create mode 100644 starbot/core/datasource.py create mode 100644 starbot/exception/DataSourceException.py diff --git a/starbot/core/bot.py b/starbot/core/bot.py new file mode 100644 index 0000000..81809ee --- /dev/null +++ b/starbot/core/bot.py @@ -0,0 +1,58 @@ +import sys + +from loguru import logger + +from starbot.core.datasource import DataSource +from starbot.exception.DataSourceException import DataSourceException + + +class StarBot: + """ + StarBot 类 + """ + STARBOT_ASCII_LOGO = "\n".join( + ( + r" _____ _ ____ _ ", + r" / ____| | | _ \ | | ", + r" | (___ | |_ __ _ _ __| |_) | ___ | |_ ", + r" \___ \| __/ _` | '__| _ < / _ \| __|", + r" ____) | || (_| | | | |_) | (_) | |_ ", + r" |_____/ \__\__,_|_| |____/ \___/ \__|", + r" StarBot - (v1.0.0) 2022-10-29", + r" Github: https://github.com/Starlwr/StarBot", + r"", + r"", + ) + ) + + def __init__(self, datasource: DataSource): + """ + Args: + datasource: 推送配置数据源 + """ + self.__datasource = datasource + + async def run(self): + """ + 启动 StarBot + """ + + # 设置日志格式 + logger_format = ( + "{time:YYYY-MM-DD HH:mm:ss.SSS} | " + "{level: <8} | " + "{name}:{line} | " + "{message}" + ) + logger.remove() + logger.add(sys.stderr, format=logger_format, level="INFO") + + logger.opt(colors=True, raw=True).info(f"{self.STARBOT_ASCII_LOGO}") + logger.info("开始启动 StarBot") + + # 从数据源中加载配置 + try: + await self.__datasource.load() + except DataSourceException as ex: + logger.error(ex.msg) + return diff --git a/starbot/core/datasource.py b/starbot/core/datasource.py new file mode 100644 index 0000000..5f44fff --- /dev/null +++ b/starbot/core/datasource.py @@ -0,0 +1,114 @@ +import abc +from typing import Union, List, Dict + +from loguru import logger +from pydantic import ValidationError + +from .model import Up, Bot +from ..exception.DataSourceException import DataSourceException + + +class DataSource(metaclass=abc.ABCMeta): + """ + Bot 推送配置数据源基类 + """ + + def __init__(self): + self.bots = {} + self.__up_list = [] + self.__uid_list = [] + self.__up_map = {} + + @abc.abstractmethod + async def load(self): + """ + 读取配置，基类空实现 + """ + pass + + def format_data(self): + """ + 处理读取后的配置 + + Raises: + DataSourceException: 配置中包含重复 uid + """ + self.__up_list = [x for up in map(lambda bot: bot.ups, self.bots.values()) for x in up] + self.__uid_list = list(map(lambda up: up.uid, self.__up_list)) + if len(set(self.__uid_list)) < len(self.__uid_list): + raise DataSourceException("配置中不可含有重复的 UID") + self.__up_map = dict(zip(map(lambda up: up.uid, self.__up_list), self.__up_list)) + + def get_up_list(self) -> List[Up]: + """ + 获取数据源中所有的 UP 实例 + + Returns: + 数据源中所有的 UP 实例列表 + """ + return self.__up_list + + def get_uid_list(self) -> List[int]: + """ + 获取数据源中所有的 UID + + Returns: + 数据源中所有的 UID 列表 + """ + return self.__uid_list + + def get_up(self, uid: int) -> Up: + """ + 根据 UID 获取 Up 实例 + + Args: + uid: 需要获取 Up 的 UID + + Returns: + Up 实例 + + Raises: + DataSourceException: uid 不存在 + """ + up = self.__up_map.get(uid) + if up is None: + raise DataSourceException(f"不存在的 UID: {uid}") + return up + + +class DictDataSource(DataSource): + """ + 使用字典结构初始化的 Bot 推送配置数据源 + """ + + def __init__(self, dict_config: Union[List[Dict], Dict]): + """ + Args: + dict_config: 配置字典 + """ + super().__init__() + logger.info("已选用 Dict 作为 Bot 数据源") + self.__config = dict_config + + if isinstance(self.__config, dict): + self.__config = [self.__config] + + async def load(self): + """ + 从配置字典中初始化配置 + + Raises: + DataSourceException: 配置字典格式错误或缺少必要参数 + """ + logger.info("开始从 Dict 中初始化 Bot 配置") + + for bot in self.__config: + if "qq" not in bot: + raise DataSourceException("提供的配置字典中未提供 Bot 的 QQ 号参数") + try: + self.bots.update({bot["qq"]: Bot(**bot)}) + except ValidationError as ex: + raise DataSourceException(f"提供的配置字典中缺少必须的 {ex.errors()[0].get('loc')[-1]} 参数") + + super().format_data() + logger.success(f"成功从 Dict 中初始化了 {len(self.get_up_list())} 个 UP 主") diff --git a/starbot/exception/DataSourceException.py b/starbot/exception/DataSourceException.py new file mode 100644 index 0000000..88a78ac --- /dev/null +++ b/starbot/exception/DataSourceException.py @@ -0,0 +1,15 @@ +""" +数据源异常 +""" + +from .ApiException import ApiException + + +class DataSourceException(ApiException): + """ + 数据源异常 + """ + + def __init__(self, msg: str): + super().__init__() + self.msg = msg diff --git a/starbot/exception/__init__.py b/starbot/exception/__init__.py index 4797575..ccffa93 100644 --- a/starbot/exception/__init__.py +++ b/starbot/exception/__init__.py @@ -2,6 +2,7 @@ from .ApiException import * from .CredentialNoBiliJctException import * from .CredentialNoBuvid3Exception import * from .CredentialNoSessdataException import * +from .DataSourceException import * from .LiveException import * from .NetworkException import * from .ResponseCodeException import *