TelegramBotAdmin/wbb/utils/cust_filter.py

import shlex
from typing import List
import re

from pyrogram.filters import create
from pyrogram.types import Message

from wbb import Command
from wbb.utils.botinfo import BOT_USERNAME


def command(
    commands: str or List[str],
    prefixes: str or List[str] = "/",
    case_sensitive: bool = False,
):
    """
    This is a drop in replacement for the default cust_filter.command that is included
    in Pyrogram. The Pyrogram one does not support /command@botname type commands,
    so this custom filter enables that throughout all groups and private chats.
    This filter works exactly the same as the original command filter even with support for multiple command
    prefixes and case sensitivity.
    Command arguments are given to user as message.command
    """

    async def func(flt, _, message: Message):
        text: str = message.text or message.caption
        message.command = None

        if not text:
            return False

        regex = "^({prefix})+\\b({regex})\\b(\\b@{bot_name}\\b)?(.*)".format(
            prefix="|".join(re.escape(x) for x in Command),
            regex="|".join(flt.commands).lower(),
            bot_name=BOT_USERNAME,
        )

        matches = re.search(re.compile(regex), text.lower())
        if matches:
            message.command = [matches.group(2)]
            try:
                for arg in shlex.split(matches.group(4).strip()):
                    if arg == BOT_USERNAME:
                        continue
                    message.command.append(arg)
            except ValueError:
                return True
            return True
        return False

    commands = commands if type(commands) is list else [commands]
    commands = {c if case_sensitive else c.lower() for c in commands}

    prefixes = [] if prefixes is None else prefixes
    prefixes = prefixes if type(prefixes) is list else [prefixes]
    prefixes = set(prefixes) if prefixes else {""}

    return create(
        func,
        "CustomCommandFilter",
        commands=commands,
        prefixes=prefixes,
        case_sensitive=case_sensitive,
    )