Allow roles to be assignable over multiple months

This includes a refactor to use a dataclass for clearer implementation.

Along with that, this changes the roles so that they're always available, but un-assignable roles are in red and give a different error.

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!",