LiteyukiBot-TriM/liteyuki/bot/__init__.py

import asyncio
import os
import platform
import sys
import threading
import time
from typing import Any, Optional

from watchdog.events import FileSystemEventHandler
from watchdog.observers import Observer

from liteyuki.bot.lifespan import LIFESPAN_FUNC, Lifespan
from liteyuki.core import IS_MAIN_PROCESS
from liteyuki.core.manager import ProcessManager
from liteyuki.core.spawn_process import mb_run, nb_run
from liteyuki.log import init_log, logger
from liteyuki.plugin import load_plugins

__all__ = ["LiteyukiBot", "get_bot"]


class LiteyukiBot:
    def __init__(self, *args, **kwargs):
        global _BOT_INSTANCE
        _BOT_INSTANCE = self  # 引用
        self.config: dict[str, Any] = kwargs
        self.init(**self.config)  # 初始化

        self.lifespan: Lifespan = Lifespan()

        self.process_manager: ProcessManager = ProcessManager(bot=self)
        self.loop = asyncio.new_event_loop()
        asyncio.set_event_loop(self.loop)
        self.loop_thread = threading.Thread(target=self.loop.run_forever, daemon=True)
        self.call_restart_count = 0

    def run(self):
        load_plugins("liteyuki/plugins")  # 加载轻雪插件

        self.loop_thread.start()  # 启动事件循环
        asyncio.run(self.lifespan.before_start())  # 启动前钩子

        self.process_manager.add_target("nonebot", nb_run, **self.config)
        self.process_manager.start("nonebot")

        self.process_manager.add_target("melobot", mb_run, **self.config)
        self.process_manager.start("melobot")

        asyncio.run(self.lifespan.after_start())  # 启动后钩子

        self.start_watcher()  # 启动文件监视器

    def start_watcher(self):
        if self.config.get("debug", False):

            src_directories = (
                "liteyuki",
                "src/liteyuki_main",
                "src/liteyuki_plugins",
                "src/nonebot_plugins",
                "src/utils",
            )
            src_excludes_extensions = ("pyc",)

            logger.debug("轻雪重载 已启用，正在加载文件修改监测……")
            restart = self.restart_process

            class CodeModifiedHandler(FileSystemEventHandler):
                """
                Handler for code file changes
                """

                def on_modified(self, event):
                    if (
                        event.src_path.endswith(src_excludes_extensions)
                        or event.is_directory
                        or "__pycache__" in event.src_path
                    ):
                        return
                    logger.info(f"文件 {event.src_path} 已修改，机器人自动重启……")
                    restart()

            code_modified_handler = CodeModifiedHandler()

            observer = Observer()
            for directory in src_directories:
                observer.schedule(code_modified_handler, directory, recursive=True)
            observer.start()

    def restart(self, delay: int = 0):
        """
        重启轻雪本体
        Returns:

        """

        if self.call_restart_count < 1:
            executable = sys.executable
            args = sys.argv
            logger.info("正在重启 尹灵温...")
            time.sleep(delay)
            if platform.system() == "Windows":
                cmd = "start"
            elif platform.system() == "Linux":
                cmd = "nohup"
            elif platform.system() == "Darwin":
                cmd = "open"
            else:
                cmd = "nohup"
            self.process_manager.terminate_all()
            # 进程退出后重启
            threading.Thread(
                target=os.system, args=(f"{cmd} {executable} {' '.join(args)}",)
            ).start()
            sys.exit(0)
        self.call_restart_count += 1

    def restart_process(self, name: Optional[str] = None):
        """
        停止轻雪
        Args:
            name: 进程名称, 默认为None, 所有进程
        Returns:

        """
        logger.info("Stopping LiteyukiBot...")

        self.loop.create_task(self.lifespan.before_shutdown())  # 重启前钩子
        self.loop.create_task(self.lifespan.before_shutdown())  # 停止前钩子

        if name:
            self.process_manager.terminate(name)
        else:
            self.process_manager.terminate_all()

    def init(self, *args, **kwargs):
        """
        初始化轻雪, 自动调用
        Returns:

        """
        self.init_config()
        self.init_logger()

    def init_logger(self):
        # 修改nonebot的日志配置
        init_log(config=self.config)

    def init_config(self):
        pass

    def on_before_start(self, func: LIFESPAN_FUNC):
        """
        注册启动前的函数
        Args:
            func:

        Returns:

        """
        return self.lifespan.on_before_start(func)

    def on_after_start(self, func: LIFESPAN_FUNC):
        """
        注册启动后的函数
        Args:
            func:

        Returns:

        """
        return self.lifespan.on_after_start(func)

    def on_before_shutdown(self, func: LIFESPAN_FUNC):
        """
        注册停止前的函数，为子进程停止时调用
        Args:
            func:

        Returns:

        """
        return self.lifespan.on_before_shutdown(func)

    def on_after_shutdown(self, func: LIFESPAN_FUNC):
        """
        注册停止后的函数：未实现
        Args:
            func:

        Returns:

        """
        return self.lifespan.on_after_shutdown(func)

    def on_before_restart(self, func: LIFESPAN_FUNC):
        """
        注册重启前的函数，为子进程重启时调用
        Args:
            func:

        Returns:

        """

        return self.lifespan.on_before_restart(func)

    def on_after_restart(self, func: LIFESPAN_FUNC):
        """
        注册重启后的函数：未实现
        Args:
            func:

        Returns:

        """
        return self.lifespan.on_after_restart(func)

    def on_after_nonebot_init(self, func: LIFESPAN_FUNC):
        """
        注册nonebot初始化后的函数
        Args:
            func:

        Returns:

        """
        return self.lifespan.on_after_nonebot_init(func)


_BOT_INSTANCE: Optional[LiteyukiBot] = None


def get_bot() -> Optional[LiteyukiBot]:
    """
    获取轻雪实例
    Returns:
        LiteyukiBot: 当前的轻雪实例
    """
    if IS_MAIN_PROCESS:
        return _BOT_INSTANCE
    else:
        # 从多进程上下文中获取
        pass