Утилиты работы с командами
This commit is contained in:
688
bot/utils/argument.py
Normal file
688
bot/utils/argument.py
Normal file
@@ -0,0 +1,688 @@
|
||||
"""
|
||||
Утилиты для работы с командами бота
|
||||
"""
|
||||
from typing import Optional, Union, Dict, List, Tuple, Set
|
||||
from dataclasses import dataclass, field
|
||||
import re
|
||||
|
||||
from aiogram.types import Message
|
||||
|
||||
from configs import settings
|
||||
|
||||
__all__ = (
|
||||
'is_command',
|
||||
'find_argument',
|
||||
'get_command',
|
||||
'parse_arguments',
|
||||
'parse_flags',
|
||||
'CommandParser',
|
||||
'ParsedCommand',
|
||||
'parse_command',
|
||||
'validate_command',
|
||||
'get_command_usage',
|
||||
'extract_mentions',
|
||||
'extract_user_ids',
|
||||
'extract_hashtags'
|
||||
)
|
||||
|
||||
|
||||
@dataclass
|
||||
class ParsedCommand:
|
||||
"""
|
||||
Распарсенная команда.
|
||||
|
||||
Attributes:
|
||||
command: Название команды
|
||||
prefix: Префикс команды
|
||||
args: Список аргументов
|
||||
raw_args: Исходная строка аргументов
|
||||
flags: Словарь флагов (--flag value)
|
||||
bot_username: Username бота (если было упоминание)
|
||||
is_group_command: Команда в группе с упоминанием бота
|
||||
"""
|
||||
command: str
|
||||
prefix: str
|
||||
args: List[str] = field(default_factory=list)
|
||||
raw_args: Optional[str] = None
|
||||
flags: Dict[str, Union[str, bool]] = field(default_factory=dict)
|
||||
bot_username: Optional[str] = None
|
||||
is_group_command: bool = False
|
||||
|
||||
@property
|
||||
def has_args(self) -> bool:
|
||||
"""Есть ли аргументы"""
|
||||
return len(self.args) > 0
|
||||
|
||||
@property
|
||||
def has_flags(self) -> bool:
|
||||
"""Есть ли флаги"""
|
||||
return len(self.flags) > 0
|
||||
|
||||
def get_arg(self, index: int, default: Optional[str] = None) -> Optional[str]:
|
||||
"""Получает аргумент по индексу"""
|
||||
return self.args[index] if index < len(self.args) else default
|
||||
|
||||
def get_flag(self, name: str, default: Optional[Union[str, bool]] = None) -> Union[str, bool, None]:
|
||||
"""Получает значение флага"""
|
||||
return self.flags.get(name, default)
|
||||
|
||||
def __repr__(self) -> str:
|
||||
return (
|
||||
f"ParsedCommand(command='{self.command}', "
|
||||
f"args={self.args}, flags={self.flags})"
|
||||
)
|
||||
|
||||
|
||||
class CommandParser:
|
||||
"""
|
||||
Парсер команд бота.
|
||||
|
||||
Возможности:
|
||||
- Поддержка нескольких префиксов
|
||||
- Парсинг аргументов
|
||||
- Парсинг флагов (--flag value, -f value)
|
||||
- Поддержка упоминаний бота (@botname)
|
||||
- Парсинг quoted аргументов ("arg with spaces")
|
||||
- Валидация команд
|
||||
|
||||
Example:
|
||||
```python
|
||||
parser = CommandParser()
|
||||
|
||||
# Парсинг команды
|
||||
parsed = parser.parse("/ban @user 7d --reason спам")
|
||||
print(parsed.command) # "ban"
|
||||
print(parsed.args) # ["@user", "7d"]
|
||||
print(parsed.flags) # {"reason": "спам"}
|
||||
```
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
prefixes: Optional[List[str]] = None,
|
||||
bot_username: Optional[str] = None
|
||||
):
|
||||
"""
|
||||
Args:
|
||||
prefixes: Список префиксов (по умолчанию из settings)
|
||||
bot_username: Username бота для проверки упоминаний
|
||||
"""
|
||||
self.prefixes = prefixes or settings.PREFIX
|
||||
self.bot_username = bot_username
|
||||
|
||||
def is_command(self, text: Optional[str]) -> bool:
|
||||
"""
|
||||
Проверяет, является ли текст командой.
|
||||
|
||||
Args:
|
||||
text: Текст для проверки
|
||||
|
||||
Returns:
|
||||
bool: True если это команда
|
||||
|
||||
Example:
|
||||
>> parser.is_command("/start")
|
||||
True
|
||||
>> parser.is_command("hello")
|
||||
False
|
||||
"""
|
||||
if not text:
|
||||
return False
|
||||
|
||||
text = text.strip()
|
||||
|
||||
# Проверяем все префиксы
|
||||
return any(text.startswith(prefix) for prefix in self.prefixes)
|
||||
|
||||
def get_command(
|
||||
self,
|
||||
text: Optional[str],
|
||||
strip_mention: bool = True
|
||||
) -> Optional[str]:
|
||||
"""
|
||||
Извлекает название команды из текста.
|
||||
|
||||
Args:
|
||||
text: Текст сообщения
|
||||
strip_mention: Убирать упоминание бота (@botname)
|
||||
|
||||
Returns:
|
||||
Optional[str]: Название команды или None
|
||||
|
||||
Example:
|
||||
>> parser.get_command("/start@mybot arg")
|
||||
'start'
|
||||
>> parser.get_command("!help")
|
||||
'help'
|
||||
"""
|
||||
if not self.is_command(text):
|
||||
return None
|
||||
|
||||
text = text.strip()
|
||||
|
||||
# Находим префикс
|
||||
prefix = next(p for p in self.prefixes if text.startswith(p))
|
||||
|
||||
# Убираем префикс
|
||||
without_prefix = text[len(prefix):]
|
||||
|
||||
# Берем первое слово
|
||||
command = without_prefix.split()[0] if without_prefix else ""
|
||||
|
||||
# Убираем упоминание бота если есть
|
||||
if strip_mention and '@' in command:
|
||||
command = command.split('@')[0]
|
||||
|
||||
return command.lower() if command else None
|
||||
|
||||
def find_argument(self, text: Optional[str]) -> Optional[str]:
|
||||
"""
|
||||
Извлекает аргументы команды (все после команды).
|
||||
|
||||
Args:
|
||||
text: Текст сообщения
|
||||
|
||||
Returns:
|
||||
Optional[str]: Аргументы или None
|
||||
|
||||
Example:
|
||||
>> parser.find_argument("/start referrer")
|
||||
'referrer'
|
||||
>> parser.find_argument("/ban @user reason text")
|
||||
'@user reason text'
|
||||
"""
|
||||
if not self.is_command(text):
|
||||
return None
|
||||
|
||||
parts = text.strip().split(maxsplit=1)
|
||||
return parts[1] if len(parts) > 1 else None
|
||||
|
||||
@staticmethod
|
||||
def parse_arguments(
|
||||
args_text: Optional[str],
|
||||
preserve_quotes: bool = False
|
||||
) -> List[str]:
|
||||
"""
|
||||
Парсит аргументы, поддерживает кавычки.
|
||||
|
||||
Args:
|
||||
args_text: Строка аргументов
|
||||
preserve_quotes: Сохранять кавычки в результате
|
||||
|
||||
Returns:
|
||||
List[str]: Список аргументов
|
||||
|
||||
Example:
|
||||
>> parser.parse_arguments('user 7d "ban reason here"')
|
||||
['user', '7d', 'ban reason here']
|
||||
"""
|
||||
if not args_text:
|
||||
return []
|
||||
|
||||
# Regex для парсинга с кавычками
|
||||
# Поддерживает: "arg with spaces" 'arg' arg
|
||||
pattern = r'''(?:[^\s"']+|"[^"]*"|'[^']*')+'''
|
||||
matches = re.findall(pattern, args_text)
|
||||
|
||||
if preserve_quotes:
|
||||
return matches
|
||||
|
||||
# Убираем кавычки
|
||||
return [m.strip('"').strip("'") for m in matches]
|
||||
|
||||
@staticmethod
|
||||
def parse_flags(
|
||||
args: List[str]
|
||||
) -> Tuple[List[str], Dict[str, Union[str, bool]]]:
|
||||
"""
|
||||
Парсит флаги из аргументов.
|
||||
|
||||
Поддерживает:
|
||||
- --flag value
|
||||
- --flag (boolean, True)
|
||||
- -f value (короткая форма)
|
||||
|
||||
Args:
|
||||
args: Список аргументов
|
||||
|
||||
Returns:
|
||||
Tuple: (аргументы_без_флагов, словарь_флагов)
|
||||
|
||||
Example:
|
||||
>> args = ['user', '--reason', 'spam', '--silent']
|
||||
>> clean_args, flags = parser.parse_flags(args)
|
||||
>> print(clean_args) # ['user']
|
||||
>> print(flags) # {'reason': 'spam', 'silent': True}
|
||||
"""
|
||||
clean_args = []
|
||||
flags = {}
|
||||
|
||||
i = 0
|
||||
while i < len(args):
|
||||
arg = args[i]
|
||||
|
||||
# Длинный флаг --flag
|
||||
if arg.startswith('--'):
|
||||
flag_name = arg[2:]
|
||||
|
||||
# Проверяем, есть ли значение
|
||||
if i + 1 < len(args) and not args[i + 1].startswith('-'):
|
||||
flags[flag_name] = args[i + 1]
|
||||
i += 2
|
||||
else:
|
||||
# Boolean флаг
|
||||
flags[flag_name] = True
|
||||
i += 1
|
||||
|
||||
# Короткий флаг -f
|
||||
elif arg.startswith('-') and len(arg) == 2:
|
||||
flag_name = arg[1]
|
||||
|
||||
# Проверяем значение
|
||||
if i + 1 < len(args) and not args[i + 1].startswith('-'):
|
||||
flags[flag_name] = args[i + 1]
|
||||
i += 2
|
||||
else:
|
||||
flags[flag_name] = True
|
||||
i += 1
|
||||
|
||||
# Обычный аргумент
|
||||
else:
|
||||
clean_args.append(arg)
|
||||
i += 1
|
||||
|
||||
return clean_args, flags
|
||||
|
||||
def parse(
|
||||
self,
|
||||
text: str,
|
||||
parse_flags: bool = True
|
||||
) -> Optional[ParsedCommand]:
|
||||
"""
|
||||
Полный парсинг команды.
|
||||
|
||||
Args:
|
||||
text: Текст команды
|
||||
parse_flags: Парсить флаги
|
||||
|
||||
Returns:
|
||||
Optional[ParsedCommand]: Распарсенная команда или None
|
||||
|
||||
Example:
|
||||
>> parsed = parser.parse('/ban @user 7d --reason "spam bot"')
|
||||
>> print(parsed.command) # 'ban'
|
||||
>> print(parsed.args) # ['@user', '7d']
|
||||
>> print(parsed.flags) # {'reason': 'spam bot'}
|
||||
"""
|
||||
if not self.is_command(text):
|
||||
return None
|
||||
|
||||
text = text.strip()
|
||||
|
||||
# Находим префикс
|
||||
prefix = next(p for p in self.prefixes if text.startswith(p))
|
||||
|
||||
# Убираем префикс
|
||||
without_prefix = text[len(prefix):]
|
||||
|
||||
# Разделяем на команду и аргументы
|
||||
parts = without_prefix.split(maxsplit=1)
|
||||
if not parts:
|
||||
return None
|
||||
|
||||
command_part = parts[0]
|
||||
raw_args = parts[1] if len(parts) > 1 else None
|
||||
|
||||
# Проверяем упоминание бота
|
||||
bot_username = None
|
||||
is_group_command = False
|
||||
|
||||
if '@' in command_part:
|
||||
cmd_parts = command_part.split('@')
|
||||
command_name = cmd_parts[0]
|
||||
bot_username = cmd_parts[1] if len(cmd_parts) > 1 else None
|
||||
is_group_command = True
|
||||
else:
|
||||
command_name = command_part
|
||||
|
||||
# Парсим аргументы
|
||||
args = self.parse_arguments(raw_args) if raw_args else []
|
||||
|
||||
# Парсим флаги
|
||||
flags = {}
|
||||
if parse_flags and args:
|
||||
args, flags = self.parse_flags(args)
|
||||
|
||||
return ParsedCommand(
|
||||
command=command_name.lower(),
|
||||
prefix=prefix,
|
||||
args=args,
|
||||
raw_args=raw_args,
|
||||
flags=flags,
|
||||
bot_username=bot_username,
|
||||
is_group_command=is_group_command
|
||||
)
|
||||
|
||||
def parse_from_message(
|
||||
self,
|
||||
message: Message,
|
||||
parse_flags: bool = True
|
||||
) -> Optional[ParsedCommand]:
|
||||
"""
|
||||
Парсит команду из объекта Message.
|
||||
|
||||
Args:
|
||||
message: Объект сообщения
|
||||
parse_flags: Парсить флаги
|
||||
|
||||
Returns:
|
||||
Optional[ParsedCommand]: Распарсенная команда
|
||||
|
||||
Example:
|
||||
>> parsed = parser.parse_from_message(message)
|
||||
>> if parsed:
|
||||
... print(f"Команда: {parsed.command}")
|
||||
"""
|
||||
if not message.text:
|
||||
return None
|
||||
|
||||
return self.parse(message.text, parse_flags=parse_flags)
|
||||
|
||||
|
||||
# Глобальный парсер
|
||||
_default_parser: Optional[CommandParser] = None
|
||||
|
||||
|
||||
def get_parser() -> CommandParser:
|
||||
"""Получает глобальный парсер команд"""
|
||||
global _default_parser
|
||||
if _default_parser is None:
|
||||
_default_parser = CommandParser()
|
||||
return _default_parser
|
||||
|
||||
|
||||
# ================= УДОБНЫЕ ФУНКЦИИ =================
|
||||
|
||||
def is_command(text: Optional[str]) -> bool:
|
||||
"""
|
||||
Проверяет, является ли текст командой.
|
||||
|
||||
Args:
|
||||
text: Текст для проверки
|
||||
|
||||
Returns:
|
||||
bool: True если это команда
|
||||
|
||||
Example:
|
||||
>> is_command("/start")
|
||||
True
|
||||
>> is_command("hello")
|
||||
False
|
||||
"""
|
||||
return get_parser().is_command(text)
|
||||
|
||||
|
||||
def find_argument(text: Optional[str]) -> Optional[str]:
|
||||
"""
|
||||
Извлекает аргументы команды.
|
||||
|
||||
Args:
|
||||
text: Текст команды
|
||||
|
||||
Returns:
|
||||
Optional[str]: Аргументы или None
|
||||
|
||||
Example:
|
||||
>> find_argument("/start referrer")
|
||||
'referrer'
|
||||
>> find_argument("/ban @user spam")
|
||||
'@user spam'
|
||||
"""
|
||||
return get_parser().find_argument(text)
|
||||
|
||||
|
||||
def get_command(text: Optional[str]) -> Optional[str]:
|
||||
"""
|
||||
Извлекает название команды.
|
||||
|
||||
Args:
|
||||
text: Текст сообщения
|
||||
|
||||
Returns:
|
||||
Optional[str]: Название команды или None
|
||||
|
||||
Example:
|
||||
>> get_command("/start@mybot")
|
||||
'start'
|
||||
>> get_command("!help")
|
||||
'help'
|
||||
"""
|
||||
return get_parser().get_command(text)
|
||||
|
||||
|
||||
def parse_arguments(args_text: Optional[str]) -> List[str]:
|
||||
"""
|
||||
Парсит аргументы команды.
|
||||
|
||||
Args:
|
||||
args_text: Строка аргументов
|
||||
|
||||
Returns:
|
||||
List[str]: Список аргументов
|
||||
|
||||
Example:
|
||||
>> parse_arguments('user 7d "ban reason"')
|
||||
['user', '7d', 'ban reason']
|
||||
"""
|
||||
return get_parser().parse_arguments(args_text)
|
||||
|
||||
|
||||
def parse_flags(args: List[str]) -> Tuple[List[str], Dict[str, Union[str, bool]]]:
|
||||
"""
|
||||
Парсит флаги из аргументов.
|
||||
|
||||
Args:
|
||||
args: Список аргументов
|
||||
|
||||
Returns:
|
||||
Tuple: (аргументы, флаги)
|
||||
|
||||
Example:
|
||||
>> args = ['user', '--reason', 'spam', '--silent']
|
||||
>> clean_args, flags = parse_flags(args)
|
||||
>> print(flags) # {'reason': 'spam', 'silent': True}
|
||||
"""
|
||||
return get_parser().parse_flags(args)
|
||||
|
||||
|
||||
def parse_command(text: str) -> Optional[ParsedCommand]:
|
||||
"""
|
||||
Полный парсинг команды.
|
||||
|
||||
Args:
|
||||
text: Текст команды
|
||||
|
||||
Returns:
|
||||
Optional[ParsedCommand]: Распарсенная команда
|
||||
|
||||
Example:
|
||||
>> parsed = parse_command('/ban @user --reason spam')
|
||||
>> print(parsed.command) # 'ban'
|
||||
>> print(parsed.args) # ['@user']
|
||||
>> print(parsed.flags) # {'reason': 'spam'}
|
||||
"""
|
||||
return get_parser().parse(text)
|
||||
|
||||
|
||||
# ================= ВАЛИДАЦИЯ КОМАНД =================
|
||||
|
||||
def validate_command(
|
||||
text: str,
|
||||
expected_command: str,
|
||||
min_args: int = 0,
|
||||
max_args: Optional[int] = None,
|
||||
required_flags: Optional[Set[str]] = None
|
||||
) -> Tuple[bool, Optional[str]]:
|
||||
"""
|
||||
Валидирует команду.
|
||||
|
||||
Args:
|
||||
text: Текст команды
|
||||
expected_command: Ожидаемая команда
|
||||
min_args: Минимальное количество аргументов
|
||||
max_args: Максимальное количество аргументов
|
||||
required_flags: Обязательные флаги
|
||||
|
||||
Returns:
|
||||
Tuple[bool, Optional[str]]: (валидна, сообщение_об_ошибке)
|
||||
|
||||
Example:
|
||||
>> valid, error = validate_command(
|
||||
... '/ban user',
|
||||
... 'ban',
|
||||
... min_args=1,
|
||||
... max_args=2
|
||||
... )
|
||||
>> if not valid:
|
||||
... print(error)
|
||||
"""
|
||||
parsed = parse_command(text)
|
||||
|
||||
if not parsed:
|
||||
return False, "Невалидная команда"
|
||||
|
||||
# Проверка команды
|
||||
if parsed.command != expected_command:
|
||||
return False, f"Ожидалась команда '{expected_command}'"
|
||||
|
||||
# Проверка количества аргументов
|
||||
arg_count = len(parsed.args)
|
||||
|
||||
if arg_count < min_args:
|
||||
return False, f"Недостаточно аргументов (минимум {min_args})"
|
||||
|
||||
if max_args is not None and arg_count > max_args:
|
||||
return False, f"Слишком много аргументов (максимум {max_args})"
|
||||
|
||||
# Проверка обязательных флагов
|
||||
if required_flags:
|
||||
missing_flags = required_flags - set(parsed.flags.keys())
|
||||
if missing_flags:
|
||||
return False, f"Отсутствуют обязательные флаги: {', '.join(missing_flags)}"
|
||||
|
||||
return True, None
|
||||
|
||||
|
||||
def get_command_usage(
|
||||
command: str,
|
||||
args: List[str],
|
||||
flags: Optional[Dict[str, str]] = None,
|
||||
description: Optional[str] = None
|
||||
) -> str:
|
||||
"""
|
||||
Формирует строку использования команды.
|
||||
|
||||
Args:
|
||||
command: Название команды
|
||||
args: Список аргументов
|
||||
flags: Словарь флагов с описанием
|
||||
description: Описание команды
|
||||
|
||||
Returns:
|
||||
str: Форматированная строка использования
|
||||
|
||||
Example:
|
||||
>> usage = get_command_usage(
|
||||
... 'ban',
|
||||
... ['<user>', '[duration]'],
|
||||
... {'reason': 'Причина бана', 'silent': 'Тихий бан'},
|
||||
... 'Банит пользователя'
|
||||
... )
|
||||
>> print(usage)
|
||||
"""
|
||||
lines = []
|
||||
|
||||
# Описание
|
||||
if description:
|
||||
lines.append(f"📝 {description}\n")
|
||||
|
||||
# Использование
|
||||
args_str = ' '.join(args)
|
||||
lines.append(f"<b>Использование:</b>")
|
||||
lines.append(f"<code>/{command} {args_str}</code>\n")
|
||||
|
||||
# Аргументы
|
||||
if args:
|
||||
lines.append("<b>Аргументы:</b>")
|
||||
for arg in args:
|
||||
# Определяем обязательность
|
||||
if arg.startswith('<') and arg.endswith('>'):
|
||||
lines.append(f"• {arg} - обязательный")
|
||||
elif arg.startswith('[') and arg.endswith(']'):
|
||||
lines.append(f"• {arg} - необязательный")
|
||||
lines.append("")
|
||||
|
||||
# Флаги
|
||||
if flags:
|
||||
lines.append("<b>Флаги:</b>")
|
||||
for flag, desc in flags.items():
|
||||
lines.append(f"• --{flag} - {desc}")
|
||||
|
||||
return '\n'.join(lines)
|
||||
|
||||
|
||||
# ================= ИЗВЛЕЧЕНИЕ УПОМИНАНИЙ =================
|
||||
|
||||
def extract_mentions(text: str) -> List[str]:
|
||||
"""
|
||||
Извлекает все упоминания (@username) из текста.
|
||||
|
||||
Args:
|
||||
text: Текст для анализа
|
||||
|
||||
Returns:
|
||||
List[str]: Список username (без @)
|
||||
|
||||
Example:
|
||||
>> extract_mentions("Бан @user1 и @user2")
|
||||
['user1', 'user2']
|
||||
"""
|
||||
pattern = r'@(\w+)'
|
||||
return re.findall(pattern, text)
|
||||
|
||||
|
||||
def extract_user_ids(text: str) -> List[int]:
|
||||
"""
|
||||
Извлекает все ID пользователей из текста.
|
||||
|
||||
Args:
|
||||
text: Текст для анализа
|
||||
|
||||
Returns:
|
||||
List[int]: Список ID
|
||||
|
||||
Example:
|
||||
>> extract_user_ids("Бан id123456789 и id987654321")
|
||||
[123456789, 987654321]
|
||||
"""
|
||||
pattern = r'id(\d+)'
|
||||
matches = re.findall(pattern, text)
|
||||
return [int(m) for m in matches]
|
||||
|
||||
|
||||
def extract_hashtags(text: str) -> List[str]:
|
||||
"""
|
||||
Извлекает все хештеги из текста.
|
||||
|
||||
Args:
|
||||
text: Текст для анализа
|
||||
|
||||
Returns:
|
||||
List[str]: Список хештегов (без #)
|
||||
|
||||
Example:
|
||||
>> extract_hashtags("Пост #важное #новости")
|
||||
['важное', 'новости']
|
||||
"""
|
||||
pattern = r'#(\w+)'
|
||||
return re.findall(pattern, text)
|
||||
Reference in New Issue
Block a user