Add commands filter and mixin for handlers

This commit is contained in:
jrootjunior 2019-12-03 16:03:46 +02:00
parent ea7b20e114
commit 6cb1f96ded
5 changed files with 157 additions and 6 deletions

View file

@ -2,7 +2,7 @@ from __future__ import annotations
from typing import Any, Optional, TypeVar 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 ...utils.token import extract_bot_id, validate_token
from ..methods import TelegramMethod from ..methods import TelegramMethod
from .session.aiohttp import AiohttpSession from .session.aiohttp import AiohttpSession
@ -11,7 +11,7 @@ from .session.base import BaseSession
T = TypeVar("T") T = TypeVar("T")
class BaseBot(ContextInstanceMixin): class BaseBot(ContextInstanceMixin, DataMixin):
def __init__(self, token: str, session: BaseSession = None, parse_mode: Optional[str] = None): def __init__(self, token: str, session: BaseSession = None, parse_mode: Optional[str] = None):
validate_token(token) validate_token(token)

View file

@ -103,6 +103,11 @@ class Bot(BaseBot):
Class where located all API methods 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 # Group: Getting updates
# Source: https://core.telegram.org/bots/api#getting-updates # Source: https://core.telegram.org/bots/api#getting-updates

View file

@ -1,14 +1,15 @@
from typing import Dict, Tuple, Union from typing import Dict, Tuple, Union
from .base import BaseFilter from .base import BaseFilter
from .command import Command, CommandObject
from .text import Text from .text import Text
__all__ = ("BUILTIN_FILTERS", "BaseFilter", "Text") __all__ = ("BUILTIN_FILTERS", "BaseFilter", "Text", "Command", "CommandObject")
BUILTIN_FILTERS: Dict[str, Union[Tuple[BaseFilter], Tuple]] = { BUILTIN_FILTERS: Dict[str, Union[Tuple[BaseFilter], Tuple]] = {
"update": (), "update": (),
"message": (Text,), "message": (Text, Command),
"edited_message": (Text,), "edited_message": (Text, Command),
"channel_post": (Text,), "channel_post": (Text,),
"edited_channel_post": (Text,), "edited_channel_post": (Text,),
"inline_query": (Text,), "inline_query": (Text,),

View file

@ -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

View file

@ -1,7 +1,9 @@
from abc import ABC from abc import ABC
from typing import Optional
from aiogram.api.types import Message 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): class MessageHandler(BaseHandler, ABC):
@ -14,3 +16,11 @@ class MessageHandler(BaseHandler, ABC):
@property @property
def chat(self): def chat(self):
return self.event.chat return self.event.chat
class MessageHandlerCommandMixin(BaseHandlerMixin):
@property
def command(self) -> Optional[CommandObject]:
if "command" in self.data:
return self.command["data"]
return None