diff --git a/aiogram/api/client/base.py b/aiogram/api/client/base.py index 85fd529b..bfc71d44 100644 --- a/aiogram/api/client/base.py +++ b/aiogram/api/client/base.py @@ -2,7 +2,7 @@ from __future__ import annotations from typing import Any, Optional, TypeVar -from ...utils.mixins import ContextInstanceMixin +from ...utils.mixins import ContextInstanceMixin, DataMixin from ...utils.token import extract_bot_id, validate_token from ..methods import TelegramMethod from .session.aiohttp import AiohttpSession @@ -11,7 +11,7 @@ from .session.base import BaseSession T = TypeVar("T") -class BaseBot(ContextInstanceMixin): +class BaseBot(ContextInstanceMixin, DataMixin): def __init__(self, token: str, session: BaseSession = None, parse_mode: Optional[str] = None): validate_token(token) diff --git a/aiogram/api/client/bot.py b/aiogram/api/client/bot.py index 5cf50b5b..7e133640 100644 --- a/aiogram/api/client/bot.py +++ b/aiogram/api/client/bot.py @@ -103,6 +103,11 @@ class Bot(BaseBot): Class where located all API methods """ + async def me(self) -> User: + if self not in self: + self[self] = await self.get_me() + return self[self] + # ============================================================================================= # Group: Getting updates # Source: https://core.telegram.org/bots/api#getting-updates diff --git a/aiogram/dispatcher/filters/__init__.py b/aiogram/dispatcher/filters/__init__.py index d0205309..4cc2c6e9 100644 --- a/aiogram/dispatcher/filters/__init__.py +++ b/aiogram/dispatcher/filters/__init__.py @@ -1,14 +1,15 @@ from typing import Dict, Tuple, Union from .base import BaseFilter +from .command import Command, CommandObject from .text import Text -__all__ = ("BUILTIN_FILTERS", "BaseFilter", "Text") +__all__ = ("BUILTIN_FILTERS", "BaseFilter", "Text", "Command", "CommandObject") BUILTIN_FILTERS: Dict[str, Union[Tuple[BaseFilter], Tuple]] = { "update": (), - "message": (Text,), - "edited_message": (Text,), + "message": (Text, Command), + "edited_message": (Text, Command), "channel_post": (Text,), "edited_channel_post": (Text,), "inline_query": (Text,), diff --git a/aiogram/dispatcher/filters/command.py b/aiogram/dispatcher/filters/command.py new file mode 100644 index 00000000..3299bb7b --- /dev/null +++ b/aiogram/dispatcher/filters/command.py @@ -0,0 +1,135 @@ +from __future__ import annotations + +import re +from dataclasses import dataclass, field +from typing import Any, AnyStr, Dict, List, Match, Optional, Pattern, Union + +from pydantic import root_validator + +from aiogram import Bot +from aiogram.api.types import Message +from aiogram.dispatcher.filters import BaseFilter + +CommandPatterType = Union[str, re.Pattern] + + +class Command(BaseFilter): + commands: List[CommandPatterType] + commands_prefix: str = "/" + commands_ignore_case: bool = False + commands_ignore_mention: bool = False + + @root_validator + def validate_constraints(cls, values: Dict[str, Any]) -> Dict[str, Any]: + if "commands" not in values: + raise ValueError("Commands required") + if not isinstance(values["commands"], list): + values["commands"] = [values["commands"]] + return values + + async def __call__(self, message: Message, bot: Bot) -> Union[bool, Dict[str, Any]]: + if not message.text: + return False + + return await self.parse_command(text=message.text, bot=bot) + + async def parse_command(self, text: str, bot: Bot) -> Union[bool, Dict[str, CommandObject]]: + """ + Extract command from the text and validate + + :param text: + :param bot: + :return: + """ + # First step: separate command with arguments + # "/command@mention arg1 arg2" -> "/command@mention", ["arg1 arg2"] + full_command, *args = text.split(maxsplit=1) + + # Separate command into valuable parts + # "/command@mention" -> "/", ("command", "@", "mention") + prefix, (command, _, mention) = full_command[0], full_command[1:].partition("@") + + # Validate prefixes + if prefix not in self.commands_prefix: + return False + + # Validate mention + if ( + mention + and not self.commands_ignore_mention + and mention.lower() != (await bot.me()).username.lower() + ): + return False + + # Validate command + for allowed_command in self.commands: + # Command can be presented as regexp pattern or raw string + # then need to validate that in different ways + if isinstance(allowed_command, Pattern): # Regexp + result = allowed_command.match(command) + if result: + return { + "command": CommandObject( + prefix=prefix, + command=command, + mention=mention, + args=args[0] if args else None, + match=result, + ) + } + + elif command == allowed_command: # String + return { + "command": CommandObject( + prefix=prefix, + command=command, + mention=mention, + args=args[0] if args else None, + match=None, + ) + } + + return False + + class Config: + arbitrary_types_allowed = True + + +@dataclass +class CommandObject: + """ + Instance of this object is always has command and it prefix. + Can be passed as keyword argument ``command`` to the handler + """ + + prefix: str = "/" + """Command prefix""" + command: str = "" + """Command without prefix and mention""" + mention: str = None + """Mention (if available)""" + args: str = field(repr=False, default=None) + """Command argument""" + match: Optional[Match[AnyStr]] = None + """Will be presented match result if the command is presented as regexp in filter""" + + @property + def mentioned(self) -> bool: + """ + This command has mention? + :return: + """ + return bool(self.mention) + + @property + def text(self) -> str: + """ + Generate original text from object + :return: + """ + line = self.prefix + self.command + if self.mentioned: + line += "@" + self.mention + if self.args: + line += " " + self.args + return line diff --git a/aiogram/dispatcher/handler/message.py b/aiogram/dispatcher/handler/message.py index 25b9df6e..2afdbf7b 100644 --- a/aiogram/dispatcher/handler/message.py +++ b/aiogram/dispatcher/handler/message.py @@ -1,7 +1,9 @@ from abc import ABC +from typing import Optional from aiogram.api.types import Message -from aiogram.dispatcher.handler.base import BaseHandler +from aiogram.dispatcher.filters import CommandObject +from aiogram.dispatcher.handler.base import BaseHandler, BaseHandlerMixin class MessageHandler(BaseHandler, ABC): @@ -14,3 +16,11 @@ class MessageHandler(BaseHandler, ABC): @property def chat(self): return self.event.chat + + +class MessageHandlerCommandMixin(BaseHandlerMixin): + @property + def command(self) -> Optional[CommandObject]: + if "command" in self.data: + return self.command["data"] + return None