From be15d5bd5b1a708d480557883ddd775371af0edb Mon Sep 17 00:00:00 2001 From: "airo.pi_" <47398145+AiroPi@users.noreply.github.com> Date: Sun, 24 Mar 2024 01:58:03 +0100 Subject: [PATCH] add doc & small improvements --- src/cogs/poll/__init__.py | 6 +-- src/cogs/poll/display.py | 72 ++++++++++++++-------------- src/cogs/poll/edit.py | 40 +++++++--------- src/cogs/poll/vote_menus.py | 8 ++-- src/core/response.py | 4 +- src/core/view_menus.py | 96 +++++++++++++++++++++++++++---------- 6 files changed, 133 insertions(+), 93 deletions(-) diff --git a/src/cogs/poll/__init__.py b/src/cogs/poll/__init__.py index 5fbfbe1..39ad71e 100644 --- a/src/cogs/poll/__init__.py +++ b/src/cogs/poll/__init__.py @@ -102,7 +102,7 @@ async def edit_poll(self, inter: Interaction, message: discord.Message) -> None: if poll.author_id != inter.user.id: raise NonSpecificError(_("You are not the author of this poll. You can't edit it.", _l=256)) await inter.response.send_message( - **(await PollDisplay.build(poll, self.bot)), + **(await PollDisplay(poll, self.bot)), view=await EditPoll(self, poll, message), ephemeral=True, ) @@ -133,7 +133,7 @@ async def on_submit(self, inter: discord.Interaction): self.poll.title = self.question.value self.poll.description = self.description.value await inter.response.send_message( - **(await PollDisplay.build(self.poll, self.bot)), + **(await PollDisplay(self.poll, self.bot)), view=await EditPoll(self.cog, self.poll, inter.message), ephemeral=True, ) @@ -173,7 +173,7 @@ async def on_submit(self, inter: discord.Interaction): self.poll.choices.append(db.PollChoice(poll_id=self.poll.id, label=self.choice3.value)) await inter.response.send_message( - **(await PollDisplay.build(self.poll, self.bot)), + **(await PollDisplay(self.poll, self.bot)), view=await EditPoll(self.cog, self.poll, inter.message), ephemeral=True, ) diff --git a/src/cogs/poll/display.py b/src/cogs/poll/display.py index c6ecc5d..aa9dda9 100644 --- a/src/cogs/poll/display.py +++ b/src/cogs/poll/display.py @@ -9,6 +9,7 @@ from core import Emojis, db from core.i18n import _ from core.response import MessageDisplay +from core.utils import AsyncInitMixin from .constants import ( BOOLEAN_INDEXES, @@ -28,48 +29,19 @@ from mybot import MyBot -class PollDisplay: - def __init__(self, poll: Poll, votes: dict[str, int] | None): - self.poll: Poll = poll - self.votes = votes +class PollDisplay(AsyncInitMixin, MessageDisplay): + async def __init__(self, poll: Poll, bot: MyBot, old_embed: Embed | None = None): + self.poll = poll + self.votes: dict[str, int] | None = await self.get_votes(bot) - @classmethod - async def build(cls, poll: Poll, bot: MyBot, old_embed: Embed | None = None) -> MessageDisplay: content = poll.description embed = discord.Embed(title=poll.title) - - votes: dict[str, int] | None - if poll.public_results is True: - async with bot.async_session.begin() as session: - stmt = ( - db.select(db.PollAnswer.value, func.count()) - .select_from(db.PollAnswer) - .where(db.PollAnswer.poll_id == poll.id) - .group_by(db.PollAnswer.value) - ) - - votes = { # noqa: C416, dict comprehension used for typing purposes - key: value - for key, value in (await session.execute(stmt)).all() # choice_id: vote_count - } - if poll.type == db.PollType.CHOICE: - # when we delete a choice from a poll, the votes are still in the db before commit - # so we need to filter them - votes = { - key: value for key, value in votes.items() if key in (str(choice.id) for choice in poll.choices) - } - else: - votes = None - - poll_display = cls(poll, votes) - - description_split: list[str] = [poll_display.build_end_date(), poll_display.build_legend()] - + description_split: list[str] = [self.build_end_date(), self.build_legend()] embed.description = "\n".join(description_split) if poll.public_results: - embed.add_field(name="\u200b", value=poll_display.build_graph()) - embed.color = poll_display.build_color() + embed.add_field(name="\u200b", value=self.build_graph()) + embed.color = self.build_color() if old_embed: embed.set_footer(text=old_embed.footer.text) @@ -77,7 +49,33 @@ async def build(cls, poll: Poll, bot: MyBot, old_embed: Embed | None = None) -> author = await bot.getch_user(poll.author_id) embed.set_footer(text=_("Poll created by {}", author.name if author else "unknown")) - return MessageDisplay(content=content, embed=embed) + MessageDisplay.__init__(self, content=content, embed=embed) + + async def get_votes(self, bot: MyBot) -> dict[str, int] | None: + if self.poll.public_results is False: + return None + + async with bot.async_session.begin() as session: + stmt = ( + db.select(db.PollAnswer.value, func.count()) + .select_from(db.PollAnswer) + .where(db.PollAnswer.poll_id == self.poll.id) + .group_by(db.PollAnswer.value) + ) + + votes = { # noqa: C416, dict comprehension used for typing purposes + key: value + for key, value in (await session.execute(stmt)).all() # choice_id: vote_count + } + if self.poll.type == db.PollType.CHOICE: + # when we delete a choice from a poll, the votes are still in the db before commit + # so we need to filter them + votes = { + key: value + for key, value in votes.items() + if key in (str(choice.id) for choice in self.poll.choices) + } + return votes @property def total_votes(self) -> int: diff --git a/src/cogs/poll/edit.py b/src/cogs/poll/edit.py index 8327ebf..f14e66f 100644 --- a/src/cogs/poll/edit.py +++ b/src/cogs/poll/edit.py @@ -52,7 +52,7 @@ async def update(self) -> None: self.toggle_poll.style = discord.ButtonStyle.red async def message_display(self) -> MessageDisplay: - return await PollDisplay.build(self.poll, self.bot) + return await PollDisplay(self.poll, self.bot) @ui.button(row=4, style=discord.ButtonStyle.red) async def reset_votes(self, inter: discord.Interaction, button: ui.Button[Self]): @@ -63,7 +63,7 @@ async def reset_votes(self, inter: discord.Interaction, button: ui.Button[Self]) async def toggle_poll(self, inter: discord.Interaction, button: ui.Button[Self]): del button # unused self.poll.closed = not self.poll.closed - await self.message_refresh(inter) + await self.edit_message(inter) @ui.button(row=4, style=discord.ButtonStyle.green) async def save(self, inter: discord.Interaction, button: ui.Button[Self]): @@ -74,7 +74,7 @@ async def save(self, inter: discord.Interaction, button: ui.Button[Self]): # channel can be other type of channels like voice, but it's ok. channel = cast(discord.TextChannel, inter.channel) message = await channel.send( - **(await PollDisplay.build(self.poll, self.bot)), view=await PollPublicMenu(self.cog, self.poll) + **(await PollDisplay(self.poll, self.bot)), view=await PollPublicMenu(self.cog, self.poll) ) self.poll.message_id = message.id @@ -91,7 +91,7 @@ async def save(self, inter: discord.Interaction, button: ui.Button[Self]): await inter.delete_original_response() await self.poll_message.edit( - **(await PollDisplay.build(self.poll, self.bot)), view=await PollPublicMenu(self.cog, self.poll) + **(await PollDisplay(self.poll, self.bot)), view=await PollPublicMenu(self.cog, self.poll) ) currents = self.cog.current_votes.pop(self.poll.id, None) @@ -137,12 +137,12 @@ def set_options(self, poll: db.Poll): # ) async def callback(self, inter: Interaction[MyBot]) -> None: # pyright: ignore [reportIncompatibleMethodOverride] - view = cast(EditPoll, self.view) + menu = cast(EditPoll, self.view) if self.values[0] == "public_results": - view.poll.public_results = not view.poll.public_results + menu.poll.public_results = not menu.poll.public_results elif self.values[0] == "users_can_change_answer": - view.poll.users_can_change_answer = not view.poll.users_can_change_answer - await view.message_refresh(inter) + menu.poll.users_can_change_answer = not menu.poll.users_can_change_answer + await menu.edit_message(inter) class EditPollMenus(ui.Select[EditPoll]): @@ -180,12 +180,8 @@ async def __init__(self, parent: EditPoll): await super().__init__(parent) self.poll = parent.poll - async def set_back(self, inter: discord.Interaction): - await self.parent.update() - await super().set_back(inter) - async def update_poll_display(self, inter: discord.Interaction, view: ui.View | None = None): - await inter.response.edit_message(**(await PollDisplay.build(self.poll, self.bot)), view=view or self) + await inter.response.edit_message(**(await PollDisplay(self.poll, self.bot)), view=view or self) class EditTitleAndDescription(EditSubmenu, ui.Modal): @@ -219,7 +215,7 @@ async def __init__(self, parent: EditPoll): async def on_submit(self, inter: discord.Interaction) -> None: self.poll.title = self.question.value self.poll.description = self.description.value - await self.set_back(inter) + await super().on_submit(inter) class EditEndingTime(EditSubmenu): @@ -269,7 +265,7 @@ async def __init__(self, parent: EditPoll): option = next(opt for opt in self.select_minutes.options if opt.value == default_minutes) option.default = True - async def cancel(self): + async def on_cancel(self): self.poll.end_date = self.old_value def set_time(self): @@ -319,7 +315,7 @@ async def __init__(self, parent: EditPoll): self.add_choice.label = _("Add a choice") self.remove_choice.label = _("Remove a choice") - async def cancel(self): + async def on_cancel(self): self.poll.choices = self.old_value async def update(self): @@ -378,7 +374,7 @@ async def choices_to_remove(self, inter: Interaction, select: ui.Select[Self]): await self.update() await self.parent.update_poll_display(inter, view=self) - async def cancel(self): + async def on_cancel(self): self.parent.poll.choices = self.old_value @@ -401,9 +397,9 @@ async def __init__(self, parent: EditPoll) -> None: @ui.select(cls=ui.Select[Self]) async def max_choices(self, inter: Interaction, select: ui.Select[Self]): self.parent.poll.max_answers = int(select.values[0]) - await self.message_refresh(inter) + await self.edit_message(inter) - async def cancel(self): + async def on_cancel(self): self.parent.poll.max_answers = self.old_value @@ -422,9 +418,9 @@ async def __init__(self, parent: EditPoll) -> None: @ui.select(cls=ui.RoleSelect[Self], max_values=25) async def allowed_roles(self, inter: Interaction, select: ui.RoleSelect[Self]): self.parent.poll.allowed_roles = [role.id for role in select.values] - await self.message_refresh(inter) + await self.edit_message(inter) - async def cancel(self): + async def on_cancel(self): self.parent.poll.allowed_roles = self.old_value @@ -445,4 +441,4 @@ async def reset(self, inter: discord.Interaction, button: ui.Button[Self]): del button # unused async with self.bot.async_session.begin() as session: await session.execute(delete(db.PollAnswer).where(db.PollAnswer.poll_id == self.parent.poll.id)) - await self.set_back(inter) + await self.set_menu(inter, self.parent) diff --git a/src/cogs/poll/vote_menus.py b/src/cogs/poll/vote_menus.py index 9398731..79716d0 100644 --- a/src/cogs/poll/vote_menus.py +++ b/src/cogs/poll/vote_menus.py @@ -14,7 +14,7 @@ from core import Menu, ResponseType, db, response_constructor from core.constants import Emojis from core.i18n import _ -from core.view_menus import SubMenu +from core.view_menus import SubMenuWithoutButtons from .constants import LEGEND_EMOJIS from .display import PollDisplay @@ -112,7 +112,7 @@ async def vote(self, inter: discord.Interaction, button: ui.Button[Self]): ) -class VoteMenu(SubMenu[PollPublicMenu]): +class VoteMenu(SubMenuWithoutButtons[PollPublicMenu]): async def __init__( self, parent: PollPublicMenu, poll: db.Poll, user_votes: Sequence[db.PollAnswer], base_inter: Interaction ): @@ -132,7 +132,7 @@ async def update_poll_display(self): try: message = cast(discord.Message, self.base_inter.message) # type: ignore old_embed = message.embeds[0] if message.embeds else None - await message.edit(**(await PollDisplay.build(self.poll, self.parent.bot, old_embed))) + await message.edit(**await PollDisplay(self.poll, self.parent.bot, old_embed)) except discord.NotFound: pass @@ -174,7 +174,7 @@ async def update(self): @ui.select(cls=ui.Select[Self]) async def choice(self, inter: Interaction, select: ui.Select[Self]): del select # unused - await self.message_refresh(inter, False) + await self.edit_message(inter, False) @ui.button(style=discord.ButtonStyle.red) async def remove_vote(self, inter: Interaction, button: ui.Button[Self]): diff --git a/src/core/response.py b/src/core/response.py index 12d6684..74b72c6 100644 --- a/src/core/response.py +++ b/src/core/response.py @@ -23,10 +23,10 @@ def __getitem__(self, key: str) -> Any: return self.__dict__[key] def __iter__(self) -> Iterator[str]: - return iter(self.__dict__) + return iter(self.__dataclass_fields__) def __len__(self) -> int: - return self.__dict__.__len__() + return self.__dataclass_fields__.__len__() class UneditedMessageDisplay(Mapping[str, Any]): diff --git a/src/core/view_menus.py b/src/core/view_menus.py index b1592f2..d351e70 100644 --- a/src/core/view_menus.py +++ b/src/core/view_menus.py @@ -23,6 +23,23 @@ class Menu(ui.View, AsyncInitMixin): + """Menus are special views that can be used to create interactive menus. + + For example, when you use the /poll command, the menu is the view that allow to chose what you want to edit. + Each edit menu (edit title, edit choices...) is a sub-menu of the main menu. + These sub-menu have by default a "Cancel" and a "Validate" button. + + When clicking Validate button, the parent menu is displayed again. + When clicking Cancel button, the method `cancel` is called and the parent menu is displayed again. + A sub-menu can have multiple sub-menus, creating a tree of menus. + + When the Validate or Cancel button is clicked, the parent's method `update` is called to update the view. + + A sub-menu can be a Modal. + + We can attach a message to a menu. This is useful to edit the view when the menu timeout. + """ + parent = None async def __init__( @@ -35,22 +52,37 @@ async def __init__( del kwargs # unused super().__init__(timeout=timeout) - self._bot = bot + self.bot = bot self.message_attached_to: discord.Message | None = message_attached_to - @property - def bot(self) -> MyBot: - return self._bot - - async def set_menu(self, inter: Interaction, menu: Menu) -> None: + async def set_menu(self, inter: Interaction, menu: Menu, view_only: bool = False) -> None: + """Set the display to a new menu.""" + await menu.update() if isinstance(menu, ui.Modal): await inter.response.send_modal(menu) else: - await inter.response.edit_message(**(await menu.message_display()), view=menu) + await inter.response.edit_message(**await menu.message_display(), view=menu) - async def update(self) -> None: + async def edit_message(self, inter: Interaction, view_only: bool = False): + """This edit the message with the view and the message content. + + It is very similar to `set_menu`, and can be considered as a "re-set" of the menu (`menu.set_menu(inter, menu)`) + + Args: + inter: the Interaction object + refresh_display: when set to False, only the view in updated, not the message content. Defaults to True. """ - Update the view with new values. Will set selected values as default for selects by default. + await self.update() + if view_only: + await inter.response.edit_message(view=self) + else: + await inter.response.edit_message(view=self, **await self.message_display()) + + async def update(self) -> None: + """Update the components to match the datas. + + The default behavior is to set the selected values as default for ui.Select children, so the user can see what + he selected. """ for item in self.children: if isinstance(item, ui.Select): @@ -58,6 +90,7 @@ async def update(self) -> None: option.default = option.value in item.values def disable_view(self): + """A utility method to disable all buttons and selects in the view.""" for item in self.children: if isinstance(item, ui.Button | ui.Select): item.disabled = True @@ -69,21 +102,31 @@ async def on_timeout(self) -> None: await self.message_attached_to.edit(view=self) async def message_display(self) -> MessageDisplay | UneditedMessageDisplay: - """This function can be defined and used in order to add a message content (embeds, etc...) within the menu.""" - return UneditedMessageDisplay() + """This function can be defined and used in order to add a message content (embeds, etc...) within the menu. - async def message_refresh(self, inter: Interaction, refresh_display: bool = True): - await self.update() - if refresh_display: - await inter.response.edit_message(view=self, **await self.message_display()) - else: - await inter.response.edit_message(view=self) + Then the message content is synchronized with the view. + """ + return UneditedMessageDisplay() @staticmethod def generate_custom_id() -> str: + """A utility method to generate a random custom id for the menu.""" return os.urandom(16).hex() +class SubMenuWithoutButtons(Menu, Generic[P]): + async def __init__( + self, + parent: P, + timeout: float | None = 600, + ): + await super().__init__( + bot=parent.bot, + timeout=timeout, + ) + self.parent: P = parent + + class SubMenu(Menu, Generic[P]): async def __init__( self, @@ -98,23 +141,23 @@ async def __init__( self.cancel_btn.label = _("Cancel") self.validate_btn.label = _("Validate") - async def set_back(self, inter: Interaction) -> None: - await self.parent.update() - await inter.response.edit_message(**(await self.parent.message_display()), view=self.parent) + async def on_cancel(self): + """Method called when the cancel button is clicked.""" - async def cancel(self): - pass + async def on_validate(self): + """Method called when the validate button is clicked.""" @ui.button(style=discord.ButtonStyle.grey, row=4) async def cancel_btn(self, inter: discord.Interaction, button: ui.Button[Self]): del button # unused - await self.cancel() - await self.set_back(inter) + await self.on_cancel() + await self.set_menu(inter, self.parent) @ui.button(style=discord.ButtonStyle.green, row=4) async def validate_btn(self, inter: discord.Interaction, button: ui.Button[Self]): del button # unused - await self.set_back(inter) + await self.on_validate() + await self.set_menu(inter, self.parent) class ModalSubMenu(Generic[P], Menu, ui.Modal): @@ -133,3 +176,6 @@ async def __init__( ) self.parent: P = parent + + async def on_submit(self, inter: discord.Interaction) -> None: + await self.set_menu(inter, self.parent)