PagerMaid-Pyro/pagermaid/hook.py

import asyncio

from pyrogram import StopPropagation

from pagermaid import hook_functions, logs
from pagermaid.single_utils import Message


class Hook:
    @staticmethod
    def on_startup():
        """
        注册一个启动钩子
        """
        def decorator(function):
            hook_functions["startup"].add(function)
            return function
        return decorator

    @staticmethod
    def on_shutdown():
        """
        注册一个关闭钩子
        """
        def decorator(function):
            hook_functions["shutdown"].add(function)
            return function
        return decorator

    @staticmethod
    def command_preprocessor():
        """
        注册一个命令预处理钩子
        """

        def decorator(function):
            hook_functions["command_pre"].add(function)
            return function

        return decorator

    @staticmethod
    def command_postprocessor():
        """
        注册一个命令后处理钩子
        """

        def decorator(function):
            hook_functions["command_post"].add(function)
            return function

        return decorator

    @staticmethod
    async def startup():
        if cors := [startup() for startup in hook_functions["startup"]]:
            try:
                await asyncio.gather(*cors)
            except Exception as exception:
                logs.info(f"[startup]: {type(exception)}: {exception}")

    @staticmethod
    async def shutdown():
        if cors := [shutdown() for shutdown in hook_functions["shutdown"]]:
            try:
                await asyncio.gather(*cors)
            except Exception as exception:
                logs.info(f"[shutdown]: {type(exception)}: {exception}")

    @staticmethod
    async def command_pre(message: Message):
        if cors := [pre(message) for pre in hook_functions["command_pre"]]:  # noqa
            try:
                await asyncio.gather(*cors)
            except StopPropagation as e:
                raise StopPropagation from e
            except Exception as exception:
                logs.info(f"[command_pre]: {type(exception)}: {exception}")

    @staticmethod
    async def command_post(message: Message):
        if cors := [post(message) for post in hook_functions["command_post"]]:  # noqa
            try:
                await asyncio.gather(*cors)
            except StopPropagation as e:
                raise StopPropagation from e
            except Exception as exception:
                logs.info(f"[command_post]: {type(exception)}: {exception}")