import asyncio from importlib import import_module from pathlib import Path from typing import Dict, Generic, List, Optional, TYPE_CHECKING, Type, TypeVar from arkowrapper import ArkoWrapper from async_timeout import timeout from typing_extensions import ParamSpec from gram_core.base_service import BaseServiceType, ComponentType, DependenceType, get_all_services from gram_core.config import config as bot_config from utils.const import PLUGIN_DIR, PROJECT_ROOT from utils.helpers import gen_pkg from utils.log import logger if TYPE_CHECKING: from gram_core.application import Application from gram_core.plugin import PluginType from gram_core.builtins.executor import Executor __all__ = ("DependenceManager", "PluginManager", "ComponentManager", "ServiceManager", "Managers") R = TypeVar("R") T = TypeVar("T") P = ParamSpec("P") def _load_module(path: Path) -> None: for pkg in gen_pkg(path): try: logger.debug('正在导入 "%s"', pkg) import_module(pkg) except Exception as e: logger.exception( '在导入 "%s" 的过程中遇到了错误 [red bold]%s[/]', pkg, type(e).__name__, exc_info=e, extra={"markup": True} ) raise SystemExit from e class Manager(Generic[T]): """生命周期控制基类""" _executor: Optional["Executor"] = None _lib: Dict[Type[T], T] = {} _application: "Optional[Application]" = None def set_application(self, application: "Application") -> None: self._application = application @property def application(self) -> "Application": if self._application is None: raise RuntimeError(f"No application was set for this {self.__class__.__name__}.") return self._application @property def executor(self) -> "Executor": """执行器""" if self._executor is None: raise RuntimeError(f"No executor was set for this {self.__class__.__name__}.") return self._executor def build_executor(self, name: str): from gram_core.builtins.executor import Executor from gram_core.builtins.dispatcher import BaseDispatcher self._executor = Executor(name, dispatcher=BaseDispatcher) self._executor.set_application(self.application) class DependenceManager(Manager[DependenceType]): """基础依赖管理""" _dependency: Dict[Type[DependenceType], DependenceType] = {} @property def dependency(self) -> List[DependenceType]: return list(self._dependency.values()) @property def dependency_map(self) -> Dict[Type[DependenceType], DependenceType]: return self._dependency async def start_dependency(self) -> None: _load_module(PROJECT_ROOT / "core/dependence") for dependence in filter(lambda x: x.is_dependence, get_all_services()): dependence: Type[DependenceType] instance: DependenceType try: if hasattr(dependence, "from_config"): # 如果有 from_config 方法 instance = dependence.from_config(bot_config) # 用 from_config 实例化服务 else: instance = await self.executor(dependence) await instance.initialize() logger.success('基础服务 "%s" 启动成功', dependence.__name__) self._lib[dependence] = instance self._dependency[dependence] = instance except Exception as e: logger.exception('基础服务 "%s" 初始化失败,BOT 将自动关闭', dependence.__name__) raise SystemExit from e async def stop_dependency(self) -> None: async def task(d): try: async with timeout(5): await d.shutdown() logger.debug('基础服务 "%s" 关闭成功', d.__class__.__name__) except asyncio.TimeoutError: logger.warning('基础服务 "%s" 关闭超时', d.__class__.__name__) except Exception as e: logger.error('基础服务 "%s" 关闭错误', d.__class__.__name__, exc_info=e) tasks = [] for dependence in self._dependency.values(): tasks.append(asyncio.create_task(task(dependence))) await asyncio.gather(*tasks) class ComponentManager(Manager[ComponentType]): """组件管理""" _components: Dict[Type[ComponentType], ComponentType] = {} @property def components(self) -> List[ComponentType]: return list(self._components.values()) @property def components_map(self) -> Dict[Type[ComponentType], ComponentType]: return self._components async def init_components(self): for path in filter( lambda x: x.is_dir() and not x.name.startswith("_"), PROJECT_ROOT.joinpath("core/services").iterdir() ): _load_module(path) components = ArkoWrapper(get_all_services()).filter(lambda x: x.is_component) retry_times = 0 max_retry_times = len(components) while components: start_len = len(components) for component in list(components): component: Type[ComponentType] instance: ComponentType try: instance = await self.executor(component) self._lib[component] = instance self._components[component] = instance components = components.remove(component) except Exception as e: # pylint: disable=W0703 logger.debug('组件 "%s" 初始化失败: [red]%s[/]', component.__name__, e, extra={"markup": True}) end_len = len(list(components)) if start_len == end_len: retry_times += 1 if retry_times == max_retry_times and components: for component in components: logger.error('组件 "%s" 初始化失败', component.__name__) raise SystemExit class ServiceManager(Manager[BaseServiceType]): """服务控制类""" _services: Dict[Type[BaseServiceType], BaseServiceType] = {} @property def services(self) -> List[BaseServiceType]: return list(self._services.values()) @property def services_map(self) -> Dict[Type[BaseServiceType], BaseServiceType]: return self._services async def _initialize_service(self, target: Type[BaseServiceType]) -> BaseServiceType: instance: BaseServiceType try: if hasattr(target, "from_config"): # 如果有 from_config 方法 instance = target.from_config(bot_config) # 用 from_config 实例化服务 else: instance = await self.executor(target) await instance.initialize() logger.success('服务 "%s" 启动成功', target.__name__) return instance except Exception as e: # pylint: disable=W0703 logger.exception('服务 "%s" 初始化失败,BOT 将自动关闭', target.__name__) raise SystemExit from e async def start_services(self) -> None: for path in filter( lambda x: x.is_dir() and not x.name.startswith("_"), PROJECT_ROOT.joinpath("core/services").iterdir() ): _load_module(path) for service in filter(lambda x: not x.is_component and not x.is_dependence, get_all_services()): # 遍历所有服务类 instance = await self._initialize_service(service) self._lib[service] = instance self._services[service] = instance async def stop_services(self) -> None: """关闭服务""" if not self._services: return async def task(s): try: async with timeout(5): await s.shutdown() logger.success('服务 "%s" 关闭成功', s.__class__.__name__) except asyncio.TimeoutError: logger.warning('服务 "%s" 关闭超时', s.__class__.__name__) except Exception as e: logger.warning('服务 "%s" 关闭失败', s.__class__.__name__, exc_info=e) logger.info("正在关闭服务") tasks = [] for service in self._services.values(): tasks.append(asyncio.create_task(task(service))) await asyncio.gather(*tasks) class PluginManager(Manager["PluginType"]): """插件管理""" _plugins: Dict[Type["PluginType"], "PluginType"] = {} @property def plugins(self) -> List["PluginType"]: """所有已经加载的插件""" return list(self._plugins.values()) @property def plugins_map(self) -> Dict[Type["PluginType"], "PluginType"]: return self._plugins async def install_plugins(self) -> None: """安装所有插件""" from gram_core.plugin import get_all_plugins for path in filter(lambda x: x.is_dir(), PLUGIN_DIR.iterdir()): _load_module(path) for plugin in get_all_plugins(): plugin: Type["PluginType"] try: instance: "PluginType" = await self.executor(plugin) except Exception as e: # pylint: disable=W0703 logger.error('插件 "%s" 初始化失败', f"{plugin.__module__}.{plugin.__name__}", exc_info=e) continue self._plugins[plugin] = instance if self._application is not None: instance.set_application(self._application) await asyncio.create_task(self.plugin_install_task(plugin, instance)) @staticmethod async def plugin_install_task(plugin: Type["PluginType"], instance: "PluginType"): try: await instance.install() logger.success('插件 "%s" 安装成功', f"{plugin.__module__}.{plugin.__name__}") except Exception as e: # pylint: disable=W0703 logger.error('插件 "%s" 安装失败', f"{plugin.__module__}.{plugin.__name__}", exc_info=e) async def uninstall_plugins(self) -> None: for plugin in self._plugins.values(): try: await plugin.uninstall() except Exception as e: # pylint: disable=W0703 logger.error('插件 "%s" 卸载失败', f"{plugin.__module__}.{plugin.__name__}", exc_info=e) class Managers(DependenceManager, ComponentManager, ServiceManager, PluginManager): """BOT 除自身外的生命周期管理类"""