command.py

from typing import (
    TYPE_CHECKING,
    Any,
    Callable,
    Coroutine,
    Optional,
    Sequence,
    Union,
)

import pyrogram
from pyrogram.filters import Filter


if TYPE_CHECKING:
    from .core import Bot

CommandFunc = Union[Callable[..., Coroutine[Any, Any, None]],
                    Callable[..., Coroutine[Any, Any, Optional[str]]]]
Decorator = Callable[[CommandFunc], CommandFunc]


def desc(_desc: str) -> Decorator:
    """Sets description on a command function."""

    def desc_decorator(func: CommandFunc) -> CommandFunc:
        setattr(func, "_cmd_description", _desc)
        return func

    return desc_decorator


def usage(_usage: str, optional: bool = False) -> Decorator:
    """Sets argument usage help on a command function."""

    def usage_decorator(func: CommandFunc) -> CommandFunc:
        setattr(func, "_cmd_usage", _usage)
        setattr(func, "_cmd_usage_optional", optional)
        return func

    return usage_decorator


def alias(*aliases: str) -> Decorator:
    """Sets aliases on a command function."""

    def alias_decorator(func: CommandFunc) -> CommandFunc:
        setattr(func, "_cmd_aliases", aliases)
        return func

    return alias_decorator


def filters(_filters: Filter) -> Decorator:
    """Sets filters on a command function."""

    def filter_decorator(func: CommandFunc) -> CommandFunc:
        setattr(func, "_cmd_filters", _filters)
        return func

    return filter_decorator


class Command:
    name: str
    desc: Optional[str]
    usage: Optional[str]
    usage_optional: bool
    aliases: Sequence[str]
    filters: Optional[Filter]
    plugin: Any
    func: CommandFunc

    def __init__(self, name: str, plugin: Any, func: CommandFunc) -> None:
        self.name = name
        self.desc = getattr(func, "_cmd_description", None)
        self.usage = getattr(func, "_cmd_usage", None)
        self.usage_optional = getattr(func, "_cmd_usage_optional", False)
        self.aliases = getattr(func, "_cmd_aliases", [])
        self.filters = getattr(func, "_cmd_filters", None)
        self.plugin = plugin
        self.func = func


# Command invocation context
class Context:
    bot: "Bot"
    msg: pyrogram.types.Message
    cmd_len: int

    response: pyrogram.types.Message
    input: str
    args: Sequence[str]

    segments: Sequence[str]
    invoker: str

    def __init__(
        self,
        bot: "Bot",
        msg: pyrogram.types.Message,
        cmd_len: int,
    ) -> None:
        self.bot = bot
        self.msg = msg
        self.cmd_len = cmd_len

        # Response message to be filled later
        self.response = None  # type: ignore
        # Single argument string
        username = self.bot.user.username
        if username in self.msg.text:
            self.input = self.msg.text[self.cmd_len + 1 + len(username) :]
        else:
            self.input = self.msg.text[self.cmd_len :]

        self.segments = self.msg.command
        self.invoker = self.segments[0]

    # Lazily resolve expensive fields
    def __getattr__(self, name: str) -> Any:
        if name == "args":
            return self._get_args()

        raise AttributeError(
            f"'{type(self).__name__}' object has no attribute '{name}'"
        )

    # Argument segments
    def _get_args(self) -> Sequence[str]:
        self.args = self.segments[1:]
        return self.args

    # Wrapper for Bot.respond()
    async def respond(
        self,
        text: str,
        *,
        mode: str = "edit",
        redact: bool = True,
        msg: Optional[pyrogram.types.Message] = None,
        **kwargs: Any,
    ) -> Optional[pyrogram.types.Message]:

        self.response = await self.bot.respond(
            msg or self.msg,
            text,
            mode=mode,
            redact=redact,
            response=self.response,
            **kwargs,
        )
        return self.response