428 lines
13 KiB
Python
428 lines
13 KiB
Python
from __future__ import annotations
|
|
|
|
from abc import ABC
|
|
from copy import deepcopy
|
|
from itertools import chain
|
|
from itertools import cycle as repeat_all
|
|
from typing import (
|
|
TYPE_CHECKING,
|
|
Any,
|
|
Generator,
|
|
Generic,
|
|
Iterable,
|
|
List,
|
|
Optional,
|
|
Type,
|
|
TypeVar,
|
|
Union,
|
|
cast,
|
|
)
|
|
|
|
from aiogram.filters.callback_data import CallbackData
|
|
from aiogram.types import (
|
|
CallbackGame,
|
|
CopyTextButton,
|
|
InlineKeyboardButton,
|
|
InlineKeyboardMarkup,
|
|
KeyboardButton,
|
|
KeyboardButtonPollType,
|
|
KeyboardButtonRequestChat,
|
|
KeyboardButtonRequestUsers,
|
|
LoginUrl,
|
|
ReplyKeyboardMarkup,
|
|
SwitchInlineQueryChosenChat,
|
|
WebAppInfo,
|
|
)
|
|
|
|
ButtonType = TypeVar("ButtonType", InlineKeyboardButton, KeyboardButton)
|
|
T = TypeVar("T")
|
|
|
|
|
|
class KeyboardBuilder(Generic[ButtonType], ABC):
|
|
"""
|
|
Generic keyboard builder that helps to adjust your markup with defined shape of lines.
|
|
|
|
Works both of InlineKeyboardMarkup and ReplyKeyboardMarkup.
|
|
"""
|
|
|
|
max_width: int = 0
|
|
min_width: int = 0
|
|
max_buttons: int = 0
|
|
|
|
def __init__(
|
|
self, button_type: Type[ButtonType], markup: Optional[List[List[ButtonType]]] = None
|
|
) -> None:
|
|
if not issubclass(button_type, (InlineKeyboardButton, KeyboardButton)):
|
|
raise ValueError(f"Button type {button_type} are not allowed here")
|
|
self._button_type: Type[ButtonType] = button_type
|
|
if markup:
|
|
self._validate_markup(markup)
|
|
else:
|
|
markup = []
|
|
self._markup: List[List[ButtonType]] = markup
|
|
|
|
@property
|
|
def buttons(self) -> Generator[ButtonType, None, None]:
|
|
"""
|
|
Get flatten set of all buttons
|
|
|
|
:return:
|
|
"""
|
|
yield from chain.from_iterable(self.export())
|
|
|
|
def _validate_button(self, button: ButtonType) -> bool:
|
|
"""
|
|
Check that button item has correct type
|
|
|
|
:param button:
|
|
:return:
|
|
"""
|
|
allowed = self._button_type
|
|
if not isinstance(button, allowed):
|
|
raise ValueError(
|
|
f"{button!r} should be type {allowed.__name__!r} not {type(button).__name__!r}"
|
|
)
|
|
return True
|
|
|
|
def _validate_buttons(self, *buttons: ButtonType) -> bool:
|
|
"""
|
|
Check that all passed button has correct type
|
|
|
|
:param buttons:
|
|
:return:
|
|
"""
|
|
return all(map(self._validate_button, buttons))
|
|
|
|
def _validate_row(self, row: List[ButtonType]) -> bool:
|
|
"""
|
|
Check that row of buttons are correct
|
|
Row can be only list of allowed button types and has length 0 <= n <= 8
|
|
|
|
:param row:
|
|
:return:
|
|
"""
|
|
if not isinstance(row, list):
|
|
raise ValueError(
|
|
f"Row {row!r} should be type 'List[{self._button_type.__name__}]' "
|
|
f"not type {type(row).__name__}"
|
|
)
|
|
if len(row) > self.max_width:
|
|
raise ValueError(f"Row {row!r} is too long (max width: {self.max_width})")
|
|
self._validate_buttons(*row)
|
|
return True
|
|
|
|
def _validate_markup(self, markup: List[List[ButtonType]]) -> bool:
|
|
"""
|
|
Check that passed markup has correct data structure
|
|
Markup is list of lists of buttons
|
|
|
|
:param markup:
|
|
:return:
|
|
"""
|
|
count = 0
|
|
if not isinstance(markup, list):
|
|
raise ValueError(
|
|
f"Markup should be type 'List[List[{self._button_type.__name__}]]' "
|
|
f"not type {type(markup).__name__!r}"
|
|
)
|
|
for row in markup:
|
|
self._validate_row(row)
|
|
count += len(row)
|
|
if count > self.max_buttons:
|
|
raise ValueError(f"Too much buttons detected Max allowed count - {self.max_buttons}")
|
|
return True
|
|
|
|
def _validate_size(self, size: Any) -> int:
|
|
"""
|
|
Validate that passed size is legit
|
|
|
|
:param size:
|
|
:return:
|
|
"""
|
|
if not isinstance(size, int):
|
|
raise ValueError("Only int sizes are allowed")
|
|
if size not in range(self.min_width, self.max_width + 1):
|
|
raise ValueError(
|
|
f"Row size {size} is not allowed, range: [{self.min_width}, {self.max_width}]"
|
|
)
|
|
return size
|
|
|
|
def export(self) -> List[List[ButtonType]]:
|
|
"""
|
|
Export configured markup as list of lists of buttons
|
|
|
|
.. code-block:: python
|
|
|
|
>>> builder = KeyboardBuilder(button_type=InlineKeyboardButton)
|
|
>>> ... # Add buttons to builder
|
|
>>> markup = InlineKeyboardMarkup(inline_keyboard=builder.export())
|
|
|
|
:return:
|
|
"""
|
|
return deepcopy(self._markup)
|
|
|
|
def add(self, *buttons: ButtonType) -> "KeyboardBuilder[ButtonType]":
|
|
"""
|
|
Add one or many buttons to markup.
|
|
|
|
:param buttons:
|
|
:return:
|
|
"""
|
|
self._validate_buttons(*buttons)
|
|
markup = self.export()
|
|
|
|
# Try to add new buttons to the end of last row if it possible
|
|
if markup and len(markup[-1]) < self.max_width:
|
|
last_row = markup[-1]
|
|
pos = self.max_width - len(last_row)
|
|
head, buttons = buttons[:pos], buttons[pos:]
|
|
last_row.extend(head)
|
|
|
|
# Separate buttons to exclusive rows with max possible row width
|
|
if self.max_width > 0:
|
|
while buttons:
|
|
row, buttons = buttons[: self.max_width], buttons[self.max_width :]
|
|
markup.append(list(row))
|
|
else:
|
|
markup.append(list(buttons))
|
|
|
|
self._markup = markup
|
|
return self
|
|
|
|
def row(
|
|
self, *buttons: ButtonType, width: Optional[int] = None
|
|
) -> "KeyboardBuilder[ButtonType]":
|
|
"""
|
|
Add row to markup
|
|
|
|
When too much buttons is passed it will be separated to many rows
|
|
|
|
:param buttons:
|
|
:param width:
|
|
:return:
|
|
"""
|
|
if width is None:
|
|
width = self.max_width
|
|
|
|
self._validate_size(width)
|
|
self._validate_buttons(*buttons)
|
|
self._markup.extend(
|
|
list(buttons[pos : pos + width]) for pos in range(0, len(buttons), width)
|
|
)
|
|
return self
|
|
|
|
def adjust(self, *sizes: int, repeat: bool = False) -> "KeyboardBuilder[ButtonType]":
|
|
"""
|
|
Adjust previously added buttons to specific row sizes.
|
|
|
|
By default, when the sum of passed sizes is lower than buttons count the last
|
|
one size will be used for tail of the markup.
|
|
If repeat=True is passed - all sizes will be cycled when available more buttons
|
|
count than all sizes
|
|
|
|
:param sizes:
|
|
:param repeat:
|
|
:return:
|
|
"""
|
|
if not sizes:
|
|
sizes = (self.max_width,)
|
|
|
|
validated_sizes = map(self._validate_size, sizes)
|
|
sizes_iter = repeat_all(validated_sizes) if repeat else repeat_last(validated_sizes)
|
|
size = next(sizes_iter)
|
|
|
|
markup = []
|
|
row: List[ButtonType] = []
|
|
for button in self.buttons:
|
|
if len(row) >= size:
|
|
markup.append(row)
|
|
size = next(sizes_iter)
|
|
row = []
|
|
row.append(button)
|
|
if row:
|
|
markup.append(row)
|
|
self._markup = markup
|
|
return self
|
|
|
|
def _button(self, **kwargs: Any) -> "KeyboardBuilder[ButtonType]":
|
|
"""
|
|
Add button to markup
|
|
|
|
:param kwargs:
|
|
:return:
|
|
"""
|
|
if isinstance(callback_data := kwargs.get("callback_data", None), CallbackData):
|
|
kwargs["callback_data"] = callback_data.pack()
|
|
button = self._button_type(**kwargs)
|
|
return self.add(button)
|
|
|
|
def as_markup(self, **kwargs: Any) -> Union[InlineKeyboardMarkup, ReplyKeyboardMarkup]:
|
|
if self._button_type is KeyboardButton:
|
|
keyboard = cast(List[List[KeyboardButton]], self.export()) # type: ignore
|
|
return ReplyKeyboardMarkup(keyboard=keyboard, **kwargs)
|
|
inline_keyboard = cast(List[List[InlineKeyboardButton]], self.export()) # type: ignore
|
|
return InlineKeyboardMarkup(inline_keyboard=inline_keyboard)
|
|
|
|
def attach(self, builder: "KeyboardBuilder[ButtonType]") -> "KeyboardBuilder[ButtonType]":
|
|
if not isinstance(builder, KeyboardBuilder):
|
|
raise ValueError(f"Only KeyboardBuilder can be attached, not {type(builder).__name__}")
|
|
if builder._button_type is not self._button_type:
|
|
raise ValueError(
|
|
f"Only builders with same button type can be attached, "
|
|
f"not {self._button_type.__name__} and {builder._button_type.__name__}"
|
|
)
|
|
self._markup.extend(builder.export())
|
|
return self
|
|
|
|
|
|
def repeat_last(items: Iterable[T]) -> Generator[T, None, None]:
|
|
items_iter = iter(items)
|
|
try:
|
|
value = next(items_iter)
|
|
except StopIteration: # pragma: no cover
|
|
# Possible case but not in place where this function is used
|
|
return
|
|
yield value
|
|
finished = False
|
|
while True:
|
|
if not finished:
|
|
try:
|
|
value = next(items_iter)
|
|
except StopIteration:
|
|
finished = True
|
|
yield value
|
|
|
|
|
|
class InlineKeyboardBuilder(KeyboardBuilder[InlineKeyboardButton]):
|
|
"""
|
|
Inline keyboard builder inherits all methods from generic builder
|
|
"""
|
|
|
|
max_width: int = 8
|
|
min_width: int = 1
|
|
max_buttons: int = 100
|
|
|
|
def button(
|
|
self,
|
|
*,
|
|
text: str,
|
|
url: Optional[str] = None,
|
|
callback_data: Optional[Union[str, CallbackData]] = None,
|
|
web_app: Optional[WebAppInfo] = None,
|
|
login_url: Optional[LoginUrl] = None,
|
|
switch_inline_query: Optional[str] = None,
|
|
switch_inline_query_current_chat: Optional[str] = None,
|
|
switch_inline_query_chosen_chat: Optional[SwitchInlineQueryChosenChat] = None,
|
|
copy_text: Optional[CopyTextButton] = None,
|
|
callback_game: Optional[CallbackGame] = None,
|
|
pay: Optional[bool] = None,
|
|
**kwargs: Any,
|
|
) -> "InlineKeyboardBuilder":
|
|
return cast(
|
|
InlineKeyboardBuilder,
|
|
self._button(
|
|
text=text,
|
|
url=url,
|
|
callback_data=callback_data,
|
|
web_app=web_app,
|
|
login_url=login_url,
|
|
switch_inline_query=switch_inline_query,
|
|
switch_inline_query_current_chat=switch_inline_query_current_chat,
|
|
switch_inline_query_chosen_chat=switch_inline_query_chosen_chat,
|
|
copy_text=copy_text,
|
|
callback_game=callback_game,
|
|
pay=pay,
|
|
**kwargs,
|
|
),
|
|
)
|
|
|
|
if TYPE_CHECKING:
|
|
|
|
def as_markup(self, **kwargs: Any) -> InlineKeyboardMarkup:
|
|
"""Construct an InlineKeyboardMarkup"""
|
|
...
|
|
|
|
def __init__(self, markup: Optional[List[List[InlineKeyboardButton]]] = None) -> None:
|
|
super().__init__(button_type=InlineKeyboardButton, markup=markup)
|
|
|
|
def copy(self: "InlineKeyboardBuilder") -> "InlineKeyboardBuilder":
|
|
"""
|
|
Make full copy of current builder with markup
|
|
|
|
:return:
|
|
"""
|
|
return InlineKeyboardBuilder(markup=self.export())
|
|
|
|
@classmethod
|
|
def from_markup(
|
|
cls: Type["InlineKeyboardBuilder"], markup: InlineKeyboardMarkup
|
|
) -> "InlineKeyboardBuilder":
|
|
"""
|
|
Create builder from existing markup
|
|
|
|
:param markup:
|
|
:return:
|
|
"""
|
|
return cls(markup=markup.inline_keyboard)
|
|
|
|
|
|
class ReplyKeyboardBuilder(KeyboardBuilder[KeyboardButton]):
|
|
"""
|
|
Reply keyboard builder inherits all methods from generic builder
|
|
"""
|
|
|
|
max_width: int = 10
|
|
min_width: int = 1
|
|
max_buttons: int = 300
|
|
|
|
def button(
|
|
self,
|
|
*,
|
|
text: str,
|
|
request_users: Optional[KeyboardButtonRequestUsers] = None,
|
|
request_chat: Optional[KeyboardButtonRequestChat] = None,
|
|
request_contact: Optional[bool] = None,
|
|
request_location: Optional[bool] = None,
|
|
request_poll: Optional[KeyboardButtonPollType] = None,
|
|
web_app: Optional[WebAppInfo] = None,
|
|
**kwargs: Any,
|
|
) -> "ReplyKeyboardBuilder":
|
|
return cast(
|
|
ReplyKeyboardBuilder,
|
|
self._button(
|
|
text=text,
|
|
request_users=request_users,
|
|
request_chat=request_chat,
|
|
request_contact=request_contact,
|
|
request_location=request_location,
|
|
request_poll=request_poll,
|
|
web_app=web_app,
|
|
**kwargs,
|
|
),
|
|
)
|
|
|
|
if TYPE_CHECKING:
|
|
|
|
def as_markup(self, **kwargs: Any) -> ReplyKeyboardMarkup: ...
|
|
|
|
def __init__(self, markup: Optional[List[List[KeyboardButton]]] = None) -> None:
|
|
super().__init__(button_type=KeyboardButton, markup=markup)
|
|
|
|
def copy(self: "ReplyKeyboardBuilder") -> "ReplyKeyboardBuilder":
|
|
"""
|
|
Make full copy of current builder with markup
|
|
|
|
:return:
|
|
"""
|
|
return ReplyKeyboardBuilder(markup=self.export())
|
|
|
|
@classmethod
|
|
def from_markup(cls, markup: ReplyKeyboardMarkup) -> "ReplyKeyboardBuilder":
|
|
"""
|
|
Create builder from existing markup
|
|
|
|
:param markup:
|
|
:return:
|
|
"""
|
|
return cls(markup=markup.keyboard)
|