Source code for telegram_bot_discussion.button.buttons_map

from typing import Dict, Type, Union


from .button import Button




class ButtonsMap:
    """`ButtonsMap` defines the accordance between an action and a `Button` class. It is needed for `ButtonsHandler` to find which `Button` was clicked."""

    _buttons: Dict[str, Type[Button]]

    def __init__(
        self,
        *buttons: Type[Button],
    ):
        self._buttons = dict()
        for button in buttons:
            self.add(button)



    def add(
        self,
        x: Type[Button],
    ) -> "ButtonsMap":
        if x.action() not in self._buttons:
            self._buttons[x.action()] = x
        else:
            raise Exception(f"Button '{x.action()}' is already added to `ButtonsMap`.")
        return self




    def search(
        self,
        x: Union[str, Type[Button]],
    ) -> Union[Type[Button], None]:
        if isinstance(x, str):
            return self._buttons.get(x, None)
        if issubclass(x, Button):
            return self._buttons.get(x.__name__, None)
        return None