aminoacid.util.commands

  1from __future__ import annotations
  2from functools import wraps
  3
  4from inspect import signature
  5from time import time
  6from typing import (
  7    TYPE_CHECKING,
  8    Any,
  9    Callable,
 10    Coroutine,
 11    List,
 12    Literal,
 13    Optional,
 14    TypeVar,
 15    Union,
 16)
 17
 18from .events import empty_cb
 19from ..exceptions import AminoBaseException, CheckFailed, CommandNotFound
 20
 21if TYPE_CHECKING:
 22    from ..abc import Context
 23
 24T = TypeVar("T")
 25
 26
 27class UserCommand:
 28    """Command defined by User"""
 29
 30    def __init__(
 31        self,
 32        func: Callable[..., Coroutine[Any, Any, T]],
 33        command_name: str = "",
 34        check: Optional[Callable[[Context], bool]] = lambda ctx: True,
 35        check_any: Optional[List[Callable[[Context], bool]]] = list(),
 36        cooldown: Optional[int] = 0,
 37    ) -> None:
 38        """Initialises a new UserCommand with a given function to call and a given name
 39
 40        Parameters
 41        ----------
 42        func : Callable[..., Coroutine[Any, Any, T]]
 43            The function that is called when the command is executed
 44        command_name : str, optional
 45            Name of the command, by default the function name
 46        check : Optional[Callable[[Context], bool]], optional
 47            Function which is called to see if the command may be called, by default always returns True
 48        check_any : Optional[List[Callable[[Context], bool]]], optional
 49            List of checks, command will execute if any of them return True, by default []
 50        cooldown : Optional[int]
 51            Time to sleep before another instance of this command can be executed by the same user
 52        """
 53
 54        self.callback = func
 55        self.check = check
 56        self.check_any = check_any
 57        self.cooldown = cooldown
 58        self.name = command_name or func.__name__
 59
 60        self.handler: Callable[
 61            [AminoBaseException, Context], Coroutine[Any, Any, T]
 62        ] = empty_cb
 63        """Handler containing a function to handle exceptions raised
 64        
 65        Parameters
 66        ----------
 67        exc : AminoBaseException
 68            The exception raised by the command
 69        ctx : Context
 70            The context of the command
 71        """
 72        self.calls = {}
 73
 74    def error(self):
 75        """Register a function to be the error handler of this command, calls `self.register_handler()`"""
 76
 77        def wrap(f: Callable[[AminoBaseException, Context], Coroutine[Any, Any, T]]):
 78            @wraps(f)
 79            def func(self: UserCommand):
 80                self.register_handler(f)
 81
 82            return func(self)
 83
 84        return wrap
 85
 86    def register_handler(
 87        self, func: Callable[[AminoBaseException, Context], Coroutine[Any, Any, T]]
 88    ):
 89        """Register a given function as a handler, this is used by the error decorator
 90
 91        Parameters
 92        ----------
 93        func : Callable[[AminoBaseException, Context]]
 94            the function to register as the handler, this is called with the exception and the command context
 95        """
 96        self.handler = func
 97
 98    def get_signature(self) -> str:
 99        """Returns the signature of the Command
100
101        Returns
102        -------
103        str
104            Signature of the command
105        """
106
107        params = signature(self.callback).parameters
108        result = list()
109        for name, param in params.items():
110            optional = bool(param.default)
111            annotation: Any = param.annotation
112            origin = getattr(annotation, "__origin__", None)
113            if origin is Union:
114                none_cls = type(None)
115                union_args = annotation.__args__
116                optional = union_args[-1] is none_cls
117                if len(union_args) == 2 and optional:
118                    annotation = union_args[0]
119                    origin = getattr(annotation, "__origin__", None)
120            if origin is Literal:
121                name = "|".join(
122                    f'"{v}"' if isinstance(v, str) else str(v)
123                    for v in annotation.__args__
124                )
125            if optional:
126                result.append(f"[{name}]")
127            elif param.kind == param.VAR_POSITIONAL:
128                result.append(f"[{name}...]")
129            else:
130                result.append(f"<{name}>")
131        return " ".join(result)
132
133    def __call__(
134        self, context: Context, *args: Any, **kwargs: Any
135    ) -> Coroutine[Any, Any, T]:
136        """Allow the Command to be executed by calling the UserCommand instance.
137        like `await UserCommand()`
138
139        Parameters
140        ----------
141        context : Context
142            Context to pass to the callback
143
144        Returns
145        -------
146        T
147            returns the Callback return value
148        """
149
150        if not any(
151            check(context) for check in self.check_any or [lambda _: True]
152        ) or not self.check(context):
153            return context.client.logger.exception(CheckFailed(context))
154        if (current_time := int(time())) <= self.calls.get(
155            context.author.id, 0
156        ) + self.cooldown:
157            return
158        self.calls[context.author.id] = current_time
159        context.client.logger.warning((self.calls, self.cooldown))
160        return self.callback(context, *args, **kwargs)
161
162    def __str__(self) -> str:
163        return self.get_signature()
164
165    def __repr__(self) -> str:
166        return str(self)
167
168
169class HelpCommand(UserCommand):
170    def __init__(self) -> None:
171        super().__init__(self.help, "help")
172
173    async def help(self, ctx: Context, command: str = ""):
174        if not command:
175            await ctx.send(
176                f"\n".join(
177                    [
178                        ctx.client.prefix + name
179                        for name, _ in ctx.client.__command_map__.items()
180                    ]
181                )
182            )
183        if command not in ctx.client.__command_map__:
184            return ctx.client.logger.exception(CommandNotFound(ctx))
185        await ctx.send(ctx.client.__command_map__[command].get_signature())