diff options
| -rw-r--r-- | bot/exts/info/subscribe.py | 95 |
1 files changed, 72 insertions, 23 deletions
diff --git a/bot/exts/info/subscribe.py b/bot/exts/info/subscribe.py index 17bb24dca..9b96e7ab2 100644 --- a/bot/exts/info/subscribe.py +++ b/bot/exts/info/subscribe.py @@ -1,3 +1,7 @@ +import calendar +import typing as t +from dataclasses import dataclass + import arrow import discord from discord.ext import commands @@ -9,15 +13,44 @@ from bot.decorators import in_whitelist from bot.log import get_logger from bot.utils import checks, members, scheduling -# Tuple of tuples, where each inner tuple is a role id and a month number. -# The month number signifies what month the role should be assignable, -# use None for the month number if it should always be active. + +@dataclass(frozen=True) +class AssignableRole: + """ + A role that can be assigned to a user. + + months_available is a tuple that signifies what months the role should be + self-assignable, using None for when it should always be available. + """ + + role_id: int + months_available: t.Optional[tuple[int]] + name: t.Optional[str] = None # This gets populated within Subscribe.init_cog() + + def is_currently_available(self) -> bool: + """Check if the role is available for the current month.""" + if self.months_available is None: + return True + return arrow.utcnow().month in self.months_available + + def get_readable_available_months(self) -> str: + """Get a readable string of the months the role is available.""" + if self.months_available is None: + return f"{self.name} is always available." + + # Join the months together with comma separators, but use "and" for the final seperator. + month_names = [calendar.month_name[month] for month in self.months_available] + available_months_str = ", ".join(month_names[:-1]) + f" and {month_names[-1]}" + return f"{self.name} can only be assigned during {available_months_str}." + + ASSIGNABLE_ROLES = ( - (constants.Roles.announcements, None), - (constants.Roles.pyweek_announcements, None), - (constants.Roles.lovefest, 2), - (constants.Roles.advent_of_code, 12), + AssignableRole(constants.Roles.announcements, None), + AssignableRole(constants.Roles.pyweek_announcements, None), + AssignableRole(constants.Roles.lovefest, (1, 2)), + AssignableRole(constants.Roles.advent_of_code, (11, 12)), ) + ITEMS_PER_ROW = 3 DELETE_MESSAGE_AFTER = 300 # Seconds @@ -47,14 +80,22 @@ class SingleRoleButton(discord.ui.Button): ADD_STYLE = discord.ButtonStyle.success REMOVE_STYLE = discord.ButtonStyle.secondary - LABEL_FORMAT = "{action} role {role_name}" + UNAVAILABLE_STYLE = discord.ButtonStyle.red + LABEL_FORMAT = "{action} role {role_name}." CUSTOM_ID_FORMAT = "subscribe-{role_id}" - def __init__(self, role: discord.Role, assigned: bool, row: int): + def __init__(self, role: AssignableRole, assigned: bool, row: int): + if role.is_currently_available(): + style = self.REMOVE_STYLE if assigned else self.ADD_STYLE + label = self.LABEL_FORMAT.format(action="Remove" if assigned else "Add", role_name=role.name) + else: + style = self.UNAVAILABLE_STYLE + label = role.name + super().__init__( - style=self.REMOVE_STYLE if assigned else self.ADD_STYLE, - label=self.LABEL_FORMAT.format(action="Remove" if assigned else "Add", role_name=role.name), - custom_id=self.CUSTOM_ID_FORMAT.format(role_id=role.id), + style=style, + label=label, + custom_id=self.CUSTOM_ID_FORMAT.format(role_id=role.role_id), row=row, ) self.role = role @@ -68,10 +109,14 @@ class SingleRoleButton(discord.ui.Button): self.view.stop() return + if not self.role.is_currently_available(): + await interaction.response.send_message(self.role.get_readable_available_months(), ephemeral=True) + return + await members.handle_role_change( interaction.user, interaction.user.remove_roles if self.assigned else interaction.user.add_roles, - self.role, + discord.Object(self.role.role_id), ) self.assigned = not self.assigned @@ -98,24 +143,27 @@ class Subscribe(commands.Cog): def __init__(self, bot: Bot): self.bot = bot self.init_task = scheduling.create_task(self.init_cog(), event_loop=self.bot.loop) - self.assignable_roles: list[discord.Role] = [] + self.assignable_roles: list[AssignableRole] = [] self.guild: discord.Guild = None async def init_cog(self) -> None: """Initialise the cog by resolving the role IDs in ASSIGNABLE_ROLES to role names.""" await self.bot.wait_until_guild_available() - current_month = arrow.utcnow().month self.guild = self.bot.get_guild(constants.Guild.id) - for role_id, month_available in ASSIGNABLE_ROLES: - if month_available is not None and month_available != current_month: - continue - role = self.guild.get_role(role_id) - if role is None: - log.warning("Could not resolve %d to a role in the guild, skipping.", role_id) + for role in ASSIGNABLE_ROLES: + discord_role = self.guild.get_role(role.role_id) + if discord_role is None: + log.warning("Could not resolve %d to a role in the guild, skipping.", role.role_id) continue - self.assignable_roles.append(role) + self.assignable_roles.append( + AssignableRole( + role_id=role.role_id, + months_available=role.months_available, + name=discord_role.name, + ) + ) @commands.cooldown(1, 10, commands.BucketType.member) @commands.command(name="subscribe") @@ -125,9 +173,10 @@ class Subscribe(commands.Cog): await self.init_task button_view = RoleButtonView(ctx.author) + author_roles = [role.id for role in ctx.author.roles] for index, role in enumerate(self.assignable_roles): row = index // ITEMS_PER_ROW - button_view.add_item(SingleRoleButton(role, role in ctx.author.roles, row)) + button_view.add_item(SingleRoleButton(role, role.role_id in author_roles, row)) await ctx.send( "Click the buttons below to add or remove your roles!", |