mirror of
https://github.com/TeamPGM/PagerMaid-Pyro.git
synced 2024-11-21 18:18:17 +00:00
56 lines
1.9 KiB
Python
56 lines
1.9 KiB
Python
from typing import NewType, Callable, Any, Awaitable, Union, TYPE_CHECKING, Optional
|
|
|
|
from ..inject import inject
|
|
|
|
if TYPE_CHECKING:
|
|
from . import Client, Message
|
|
|
|
CommandHandlerFunc = NewType("CommandHandlerFunc", Callable[[Any, Any], Awaitable[Any]])
|
|
CommandHandlerDecorator = NewType(
|
|
"CommandHandlerDecorator",
|
|
Callable[[Union["CommandHandler", CommandHandlerFunc]], "CommandHandler"],
|
|
)
|
|
|
|
|
|
class CommandHandler:
|
|
def __init__(self, func: CommandHandlerFunc, command: Optional[str]) -> None:
|
|
self._pgp_func__: CommandHandlerFunc = func
|
|
self._pgp_command__: Optional[str] = command
|
|
self._pgp_raw_handler = None
|
|
|
|
def func(self) -> CommandHandlerFunc:
|
|
return self._pgp_func__
|
|
|
|
def __call__(self, *args, **kwargs):
|
|
return self.func()(*args, **kwargs)
|
|
|
|
def set_handler(self, handler):
|
|
self._pgp_raw_handler = handler
|
|
|
|
def get_handler(self):
|
|
return self._pgp_raw_handler
|
|
|
|
async def handler(self, client: "Client", message: "Message"):
|
|
func = self.func()
|
|
if data := inject(message, func):
|
|
await func(**data)
|
|
else:
|
|
if func.__code__.co_argcount == 0:
|
|
await func()
|
|
if func.__code__.co_argcount == 1:
|
|
await func(message)
|
|
elif func.__code__.co_argcount == 2:
|
|
await func(client, message)
|
|
|
|
def sub_command(self, **kwargs) -> CommandHandlerDecorator:
|
|
if self._pgp_command__ is None:
|
|
raise ValueError("Cannot add subcommand to a handler without a command")
|
|
if self._pgp_raw_handler is None:
|
|
raise ValueError("Cannot add subcommand to a handler without init")
|
|
from pagermaid.listener import listener
|
|
|
|
def decorator(func: CommandHandlerFunc) -> CommandHandlerFunc:
|
|
return listener(__parent_command=self._pgp_command__, **kwargs)(func)
|
|
|
|
return decorator
|