Skip to content

Commit

Permalink
Update KeyboardBuilder utility, fixed type-hints for button method, a… (
Browse files Browse the repository at this point in the history 
#1399)

* Update KeyboardBuilder utility, fixed type-hints for button method, adjusted limits of the different markup types to real world values.

* Added changelog

* Fixed coverage

* Update aiogram/utils/keyboard.py

Co-authored-by: Suren Khorenyan <surenkhorenyan@gmail.com>

* Fixed codestyle

---------

Co-authored-by: Suren Khorenyan <surenkhorenyan@gmail.com>
  • Loading branch information
@JrooTJunior @mahenzon
JrooTJunior and mahenzon committed Jan 27, 2024
1 parent d3c6379 commit 844d6f5
Show file tree
Hide file tree
Showing 3 changed files with 108 additions and 76 deletions.
1 change: 1 addition & 0 deletions CHANGES/1399.bugfix.rst
View file
Original file line number Diff line number Diff line change
@@ -0,0 +1 @@
Update KeyboardBuilder utility, fixed type-hints for button method, adjusted limits of the different markup types to real world values.
142 changes: 85 additions & 57 deletions aiogram/utils/keyboard.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -16,7 +16,6 @@
TypeVar,
Union,
cast,
no_type_check,
)

from aiogram.filters.callback_data import CallbackData
Expand All @@ -26,6 +25,8 @@
InlineKeyboardMarkup,
KeyboardButton,
KeyboardButtonPollType,
KeyboardButtonRequestChat,
KeyboardButtonRequestUsers,
LoginUrl,
ReplyKeyboardMarkup,
SwitchInlineQueryChosenChat,
Expand All @@ -34,9 +35,6 @@

ButtonType = TypeVar("ButtonType", InlineKeyboardButton, KeyboardButton)
T = TypeVar("T")
MAX_WIDTH = 8
MIN_WIDTH = 1
MAX_BUTTONS = 100


class KeyboardBuilder(Generic[ButtonType], ABC):
Expand All @@ -46,6 +44,10 @@ class KeyboardBuilder(Generic[ButtonType], ABC):
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:
Expand Down Expand Up @@ -103,8 +105,8 @@ def _validate_row(self, row: List[ButtonType]) -> bool:
f"Row {row!r} should be type 'List[{self._button_type.__name__}]' "
f"not type {type(row).__name__}"
)
if len(row) > MAX_WIDTH:
raise ValueError(f"Row {row!r} is too long (MAX_WIDTH={MAX_WIDTH})")
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

Expand All @@ -125,8 +127,8 @@ def _validate_markup(self, markup: List[List[ButtonType]]) -> bool:
for row in markup:
self._validate_row(row)
count += len(row)
if count > MAX_BUTTONS:
raise ValueError(f"Too much buttons detected Max allowed count - {MAX_BUTTONS}")
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:
Expand All @@ -138,18 +140,12 @@ def _validate_size(self, size: Any) -> int:
"""
if not isinstance(size, int):
raise ValueError("Only int sizes are allowed")
if size not in range(MIN_WIDTH, MAX_WIDTH + 1):
raise ValueError(f"Row size {size} are not 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 copy(self: "KeyboardBuilder[ButtonType]") -> "KeyboardBuilder[ButtonType]":
"""
Make full copy of current builder with markup
:return:
"""
return self.__class__(self._button_type, markup=self.export())

def export(self) -> List[List[ButtonType]]:
"""
Export configured markup as list of lists of buttons
Expand All @@ -175,21 +171,23 @@ def add(self, *buttons: ButtonType) -> "KeyboardBuilder[ButtonType]":
markup = self.export()

# Try to add new buttons to the end of last row if it possible
if markup and len(markup[-1]) < MAX_WIDTH:
if markup and len(markup[-1]) < self.max_width:
last_row = markup[-1]
pos = MAX_WIDTH - len(last_row)
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
while buttons:
row, buttons = buttons[:MAX_WIDTH], buttons[MAX_WIDTH:]
row, buttons = buttons[: self.max_width], buttons[self.max_width :]
markup.append(list(row))

self._markup = markup
return self

def row(self, *buttons: ButtonType, width: int = MAX_WIDTH) -> "KeyboardBuilder[ButtonType]":
def row(
self, *buttons: ButtonType, width: Optional[int] = None
) -> "KeyboardBuilder[ButtonType]":
"""
Add row to markup
Expand All @@ -199,6 +197,9 @@ def row(self, *buttons: ButtonType, width: int = MAX_WIDTH) -> "KeyboardBuilder[
:param width:
:return:
"""
if width is None:
width = self.max_width

self._validate_size(width)
self._validate_buttons(*buttons)
self._markup.extend(
Expand All @@ -220,7 +221,7 @@ def adjust(self, *sizes: int, repeat: bool = False) -> "KeyboardBuilder[ButtonTy
:return:
"""
if not sizes:
sizes = (MAX_WIDTH,)
sizes = (self.max_width,)

validated_sizes = map(self._validate_size, sizes)
sizes_iter = repeat_all(validated_sizes) if repeat else repeat_last(validated_sizes)
Expand All @@ -239,7 +240,7 @@ def adjust(self, *sizes: int, repeat: bool = False) -> "KeyboardBuilder[ButtonTy
self._markup = markup
return self

def button(self, **kwargs: Any) -> "KeyboardBuilder[ButtonType]":
def _button(self, **kwargs: Any) -> "KeyboardBuilder[ButtonType]":
"""
Add button to markup
Expand Down Expand Up @@ -293,25 +294,40 @@ class InlineKeyboardBuilder(KeyboardBuilder[InlineKeyboardButton]):
Inline keyboard builder inherits all methods from generic builder
"""

if TYPE_CHECKING:
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,
callback_game: Optional[CallbackGame] = None,
pay: Optional[bool] = None,
**kwargs: Any,
) -> "KeyboardBuilder[InlineKeyboardButton]":
return 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,
callback_game=callback_game,
pay=pay,
**kwargs,
)

@no_type_check
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,
callback_game: Optional[CallbackGame] = None,
pay: Optional[bool] = None,
**kwargs: Any,
) -> "KeyboardBuilder[InlineKeyboardButton]":
...
if TYPE_CHECKING:

def as_markup(self, **kwargs: Any) -> InlineKeyboardMarkup:
"""Construct an InlineKeyboardMarkup"""
Expand Down Expand Up @@ -346,22 +362,34 @@ class ReplyKeyboardBuilder(KeyboardBuilder[KeyboardButton]):
Reply keyboard builder inherits all methods from generic builder
"""

if TYPE_CHECKING:
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,
) -> "KeyboardBuilder[KeyboardButton]":
return 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,
)

@no_type_check
def button(
self,
*,
text: str,
request_user: Optional[bool] = None,
request_chat: Optional[bool] = None,
request_contact: Optional[bool] = None,
request_location: Optional[bool] = None,
request_poll: Optional[KeyboardButtonPollType] = None,
web_app: Optional[WebAppInfo] = None,
**kwargs: Any,
) -> "KeyboardBuilder[KeyboardButton]":
...
if TYPE_CHECKING:

def as_markup(self, **kwargs: Any) -> ReplyKeyboardMarkup:
...
Expand Down
41 changes: 22 additions & 19 deletions tests/test_utils/test_keyboard.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -61,36 +61,38 @@ def test_validate_row(self):

with pytest.raises(ValueError):
assert builder._validate_row(
row=(KeyboardButton(text=f"test {index}") for index in range(10))
row=(
KeyboardButton(text=f"test {index}") for index in range(builder.max_width + 5)
)
)

with pytest.raises(ValueError):
assert builder._validate_row(
row=[KeyboardButton(text=f"test {index}") for index in range(10)]
row=[
KeyboardButton(text=f"test {index}") for index in range(builder.max_width + 5)
]
)

for count in range(9):
for count in range(11):
assert builder._validate_row(
row=[KeyboardButton(text=f"test {index}") for index in range(count)]
)

def test_validate_markup(self):
def test_validate_markup_invalid_type(self):
builder = ReplyKeyboardBuilder()

with pytest.raises(ValueError):
builder._validate_markup(markup=())

def test_validate_markup_too_many_buttons(self):
builder = ReplyKeyboardBuilder()
with pytest.raises(ValueError):
builder._validate_markup(
markup=[
[KeyboardButton(text=f"{row}.{col}") for col in range(8)] for row in range(15)
[KeyboardButton(text=f"{row}.{col}") for col in range(builder.max_width)]
for row in range(builder.max_buttons)
]
)

assert builder._validate_markup(
markup=[[KeyboardButton(text=f"{row}.{col}") for col in range(8)] for row in range(8)]
)

def test_validate_size(self):
builder = ReplyKeyboardBuilder()
with pytest.raises(ValueError):
Expand All @@ -102,7 +104,7 @@ def test_validate_size(self):
builder._validate_size(0)

with pytest.raises(ValueError):
builder._validate_size(10)
builder._validate_size(builder.max_width + 5)
for size in range(1, 9):
builder._validate_size(size)

Expand All @@ -126,12 +128,6 @@ def test_export(self):
InlineKeyboardBuilder(markup=[[InlineKeyboardButton(text="test")]]),
InlineKeyboardButton(text="test2"),
],
[
KeyboardBuilder(
button_type=InlineKeyboardButton, markup=[[InlineKeyboardButton(text="test")]]
),
InlineKeyboardButton(text="test2"),
],
],
)
def test_copy(self, builder, button):
Expand All @@ -153,7 +149,14 @@ def test_copy(self, builder, button):

@pytest.mark.parametrize(
"count,rows,last_columns",
[[0, 0, 0], [3, 1, 3], [8, 1, 8], [9, 2, 1], [16, 2, 8], [19, 3, 3]],
[
[0, 0, 0],
[3, 1, 3],
[8, 1, 8],
[12, 2, 2],
[16, 2, 6],
[22, 3, 2],
],
)
def test_add(self, count: int, rows: int, last_columns: int):
builder = ReplyKeyboardBuilder()
Expand Down Expand Up @@ -182,8 +185,8 @@ def test_row(
[0, False, [2], []],
[1, False, [2], [1]],
[3, False, [2], [2, 1]],
[10, False, [], [8, 2]],
[10, False, [3, 2, 1], [3, 2, 1, 1, 1, 1, 1]],
[12, False, [], [10, 2]],
[12, True, [3, 2, 1], [3, 2, 1, 3, 2, 1]],
],
)
Expand Down

0 comments on commit 844d6f5

Please sign in to comment.