Merge pull request #1868 from python-discord/subscribe-with-buttons

Subscribe with buttons

6 files changed, 248 insertions, 88 deletions

diff --git a/bot/constants.py b/bot/constants.py
index 93da6a906..3170c2915 100644
--- a/bot/constants.py
+++ b/bot/constants.py
@@ -484,7 +484,12 @@ class Roles(metaclass=YAMLGetter):
     section = "guild"
     subsection = "roles"
 
+    # Self-assignable roles, see the Subscribe cog
+    advent_of_code: int
     announcements: int
+    lovefest: int
+    pyweek_announcements: int
+
     contributors: int
     help_cooldown: int
     muted: int
diff --git a/bot/exts/help_channels/_cog.py b/bot/exts/help_channels/_cog.py
index 0c411df04..60209ba6e 100644
--- a/bot/exts/help_channels/_cog.py
+++ b/bot/exts/help_channels/_cog.py
@@ -66,6 +66,9 @@ class HelpChannels(commands.Cog):
         self.bot = bot
         self.scheduler = scheduling.Scheduler(self.__class__.__name__)
 
+        self.guild: discord.Guild = None
+        self.cooldown_role: discord.Role = None
+
         # Categories
         self.available_category: discord.CategoryChannel = None
         self.in_use_category: discord.CategoryChannel = None
@@ -95,24 +98,6 @@ class HelpChannels(commands.Cog):
 
         self.scheduler.cancel_all()
 
-    async def _handle_role_change(self, member: discord.Member, coro: t.Callable[..., t.Coroutine]) -> None:
-        """
-        Change `member`'s cooldown role via awaiting `coro` and handle errors.
-
-        `coro` is intended to be `discord.Member.add_roles` or `discord.Member.remove_roles`.
-        """
-        try:
-            await coro(self.bot.get_guild(constants.Guild.id).get_role(constants.Roles.help_cooldown))
-        except discord.NotFound:
-            log.debug(f"Failed to change role for {member} ({member.id}): member not found")
-        except discord.Forbidden:
-            log.debug(
-                f"Forbidden to change role for {member} ({member.id}); "
-                f"possibly due to role hierarchy"
-            )
-        except discord.HTTPException as e:
-            log.error(f"Failed to change role for {member} ({member.id}): {e.status} {e.code}")
-
     @lock.lock_arg(NAMESPACE, "message", attrgetter("channel.id"))
     @lock.lock_arg(NAMESPACE, "message", attrgetter("author.id"))
     @lock.lock_arg(f"{NAMESPACE}.unclaim", "message", attrgetter("author.id"), wait=True)
@@ -130,7 +115,7 @@ class HelpChannels(commands.Cog):
         if not isinstance(message.author, discord.Member):
             log.debug(f"{message.author} ({message.author.id}) isn't a member. Not giving cooldown role or sending DM.")
         else:
-            await self._handle_role_change(message.author, message.author.add_roles)
+            await members.handle_role_change(message.author, message.author.add_roles, self.cooldown_role)
 
             try:
                 await _message.dm_on_open(message)
@@ -302,6 +287,9 @@ class HelpChannels(commands.Cog):
         await self.bot.wait_until_guild_available()
 
         log.trace("Initialising the cog.")
+        self.guild = self.bot.get_guild(constants.Guild.id)
+        self.cooldown_role = self.guild.get_role(constants.Roles.help_cooldown)
+
         await self.init_categories()
 
         self.channel_queue = self.create_channel_queue()
@@ -445,11 +433,11 @@ class HelpChannels(commands.Cog):
         await _caches.claimants.delete(channel.id)
         await _caches.session_participants.delete(channel.id)
 
-        claimant = await members.get_or_fetch_member(self.bot.get_guild(constants.Guild.id), claimant_id)
+        claimant = await members.get_or_fetch_member(self.guild, claimant_id)
         if claimant is None:
             log.info(f"{claimant_id} left the guild during their help session; the cooldown role won't be removed")
         else:
-            await self._handle_role_change(claimant, claimant.remove_roles)
+            await members.handle_role_change(claimant, claimant.remove_roles, self.cooldown_role)
 
         await _message.unpin(channel)
         await _stats.report_complete_session(channel.id, closed_on)
diff --git a/bot/exts/info/subscribe.py b/bot/exts/info/subscribe.py
new file mode 100644
index 000000000..2e6101d27
--- /dev/null
+++ b/bot/exts/info/subscribe.py
@@ -0,0 +1,202 @@
+import calendar
+import operator
+import typing as t
+from dataclasses import dataclass
+
+import arrow
+import discord
+from discord.ext import commands
+from discord.interactions import Interaction
+
+from bot import constants
+from bot.bot import Bot
+from bot.decorators import in_whitelist
+from bot.log import get_logger
+from bot.utils import checks, members, scheduling
+
+
+@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 = (
+    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
+
+log = get_logger(__name__)
+
+
+class RoleButtonView(discord.ui.View):
+    """A list of SingleRoleButtons to show to the member."""
+
+    def __init__(self, member: discord.Member):
+        super().__init__()
+        self.interaction_owner = member
+
+    async def interaction_check(self, interaction: Interaction) -> bool:
+        """Ensure that the user clicking the button is the member who invoked the command."""
+        if interaction.user != self.interaction_owner:
+            await interaction.response.send_message(
+                ":x: This is not your command to react to!",
+                ephemeral=True
+            )
+            return False
+        return True
+
+
+class SingleRoleButton(discord.ui.Button):
+    """A button that adds or removes a role from the member depending on it's current state."""
+
+    ADD_STYLE = discord.ButtonStyle.success
+    REMOVE_STYLE = discord.ButtonStyle.red
+    UNAVAILABLE_STYLE = discord.ButtonStyle.secondary
+    LABEL_FORMAT = "{action} role {role_name}."
+    CUSTOM_ID_FORMAT = "subscribe-{role_id}"
+
+    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 = f"🔒 {role.name}"
+
+        super().__init__(
+            style=style,
+            label=label,
+            custom_id=self.CUSTOM_ID_FORMAT.format(role_id=role.role_id),
+            row=row,
+        )
+        self.role = role
+        self.assigned = assigned
+
+    async def callback(self, interaction: Interaction) -> None:
+        """Update the member's role and change button text to reflect current text."""
+        if isinstance(interaction.user, discord.User):
+            log.trace("User %s is not a member", interaction.user)
+            await interaction.message.delete()
+            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,
+            discord.Object(self.role.role_id),
+        )
+
+        self.assigned = not self.assigned
+        await self.update_view(interaction)
+        await interaction.response.send_message(
+            self.LABEL_FORMAT.format(action="Added" if self.assigned else "Removed", role_name=self.role.name),
+            ephemeral=True,
+        )
+
+    async def update_view(self, interaction: Interaction) -> None:
+        """Updates the original interaction message with a new view object with the updated buttons."""
+        self.style = self.REMOVE_STYLE if self.assigned else self.ADD_STYLE
+        self.label = self.LABEL_FORMAT.format(action="Remove" if self.assigned else "Add", role_name=self.role.name)
+        try:
+            await interaction.message.edit(view=self.view)
+        except discord.NotFound:
+            log.debug("Subscribe message for %s removed before buttons could be updated", interaction.user)
+            self.view.stop()
+
+
+class Subscribe(commands.Cog):
+    """Cog to allow user to self-assign & remove the roles present in ASSIGNABLE_ROLES."""
+
+    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[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()
+
+        self.guild = self.bot.get_guild(constants.Guild.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(
+                AssignableRole(
+                    role_id=role.role_id,
+                    months_available=role.months_available,
+                    name=discord_role.name,
+                )
+            )
+        # Sort unavailable roles to the end of the list
+        self.assignable_roles.sort(key=operator.methodcaller("is_currently_available"), reverse=True)
+
+    @commands.cooldown(1, 10, commands.BucketType.member)
+    @commands.command(name="subscribe")
+    @in_whitelist(channels=(constants.Channels.bot_commands,))
+    async def subscribe_command(self, ctx: commands.Context, *_) -> None:  # We don't actually care about the args
+        """Display the member's current state for each role, and allow them to add/remove the roles."""
+        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.role_id in author_roles, row))
+
+        await ctx.reply(
+            "Click the buttons below to add or remove your roles!",
+            view=button_view,
+            delete_after=DELETE_MESSAGE_AFTER,
+        )
+
+    # This cannot be static (must have a __func__ attribute).
+    async def cog_command_error(self, ctx: commands.Context, error: Exception) -> None:
+        """Check for & ignore any InWhitelistCheckFailure."""
+        if isinstance(error, checks.InWhitelistCheckFailure):
+            error.handled = True
+
+
+def setup(bot: Bot) -> None:
+    """Load the Subscribe cog."""
+    if len(ASSIGNABLE_ROLES) > ITEMS_PER_ROW*5:  # Discord limits views to 5 rows of buttons.
+        log.error("Too many roles for 5 rows, not loading the Subscribe cog.")
+    else:
+        bot.add_cog(Subscribe(bot))
diff --git a/bot/exts/moderation/verification.py b/bot/exts/moderation/verification.py
index ed5571d2a..37338d19c 100644
--- a/bot/exts/moderation/verification.py
+++ b/bot/exts/moderation/verification.py
@@ -5,9 +5,7 @@ from discord.ext.commands import Cog, Context, command, has_any_role
 
 from bot import constants
 from bot.bot import Bot
-from bot.decorators import in_whitelist
 from bot.log import get_logger
-from bot.utils.checks import InWhitelistCheckFailure
 
 log = get_logger(__name__)
 
@@ -29,11 +27,11 @@ You can find a copy of our rules for reference at <https://pythondiscord.com/pag
 
 Additionally, if you'd like to receive notifications for the announcements \
 we post in <#{constants.Channels.announcements}>
-from time to time, you can send `!subscribe` to <#{constants.Channels.bot_commands}> at any time \
+from time to time, you can send `{constants.Bot.prefix}subscribe` to <#{constants.Channels.bot_commands}> at any time \
 to assign yourself the **Announcements** role. We'll mention this role every time we make an announcement.
 
-If you'd like to unsubscribe from the announcement notifications, simply send `!unsubscribe` to \
-<#{constants.Channels.bot_commands}>.
+If you'd like to unsubscribe from the announcement notifications, simply send `{constants.Bot.prefix}subscribe` to \
+<#{constants.Channels.bot_commands}> and click the role again!.
 
 To introduce you to our community, we've made the following video:
 https://youtu.be/ZH26PuX3re0
@@ -61,11 +59,9 @@ async def safe_dm(coro: t.Coroutine) -> None:
 
 class Verification(Cog):
     """
-    User verification and role management.
+    User verification.
 
     Statistics are collected in the 'verification.' namespace.
-
-    Additionally, this cog offers the !subscribe and !unsubscribe commands,
     """
 
     def __init__(self, bot: Bot) -> None:
@@ -108,67 +104,8 @@ class Verification(Cog):
                 log.exception("DM dispatch failed on unexpected error code")
 
     # endregion
-    # region: subscribe commands
-
-    @command(name='subscribe')
-    @in_whitelist(channels=(constants.Channels.bot_commands,))
-    async def subscribe_command(self, ctx: Context, *_) -> None:  # We don't actually care about the args
-        """Subscribe to announcement notifications by assigning yourself the role."""
-        has_role = False
-
-        for role in ctx.author.roles:
-            if role.id == constants.Roles.announcements:
-                has_role = True
-                break
-
-        if has_role:
-            await ctx.send(f"{ctx.author.mention} You're already subscribed!")
-            return
-
-        log.debug(f"{ctx.author} called !subscribe. Assigning the 'Announcements' role.")
-        await ctx.author.add_roles(discord.Object(constants.Roles.announcements), reason="Subscribed to announcements")
-
-        log.trace(f"Deleting the message posted by {ctx.author}.")
-
-        await ctx.send(
-            f"{ctx.author.mention} Subscribed to <#{constants.Channels.announcements}> notifications.",
-        )
-
-    @command(name='unsubscribe')
-    @in_whitelist(channels=(constants.Channels.bot_commands,))
-    async def unsubscribe_command(self, ctx: Context, *_) -> None:  # We don't actually care about the args
-        """Unsubscribe from announcement notifications by removing the role from yourself."""
-        has_role = False
-
-        for role in ctx.author.roles:
-            if role.id == constants.Roles.announcements:
-                has_role = True
-                break
-
-        if not has_role:
-            await ctx.send(f"{ctx.author.mention} You're already unsubscribed!")
-            return
-
-        log.debug(f"{ctx.author} called !unsubscribe. Removing the 'Announcements' role.")
-        await ctx.author.remove_roles(
-            discord.Object(constants.Roles.announcements), reason="Unsubscribed from announcements"
-        )
-
-        log.trace(f"Deleting the message posted by {ctx.author}.")
-
-        await ctx.send(
-            f"{ctx.author.mention} Unsubscribed from <#{constants.Channels.announcements}> notifications."
-        )
-
-    # endregion
     # region: miscellaneous
 
-    # This cannot be static (must have a __func__ attribute).
-    async def cog_command_error(self, ctx: Context, error: Exception) -> None:
-        """Check for & ignore any InWhitelistCheckFailure."""
-        if isinstance(error, InWhitelistCheckFailure):
-            error.handled = True
-
     @command(name='verify')
     @has_any_role(*constants.MODERATION_ROLES)
     async def perform_manual_verification(self, ctx: Context, user: discord.Member) -> None:
diff --git a/bot/utils/members.py b/bot/utils/members.py
index 77ddf1696..693286045 100644
--- a/bot/utils/members.py
+++ b/bot/utils/members.py
@@ -23,3 +23,26 @@ async def get_or_fetch_member(guild: discord.Guild, member_id: int) -> t.Optiona
             return None
         log.trace("%s fetched from API.", member)
     return member
+
+
+async def handle_role_change(
+    member: discord.Member,
+    coro: t.Callable[..., t.Coroutine],
+    role: discord.Role
+) -> None:
+    """
+    Change `member`'s cooldown role via awaiting `coro` and handle errors.
+
+    `coro` is intended to be `discord.Member.add_roles` or `discord.Member.remove_roles`.
+    """
+    try:
+        await coro(role)
+    except discord.NotFound:
+        log.debug(f"Failed to change role for {member} ({member.id}): member not found")
+    except discord.Forbidden:
+        log.debug(
+            f"Forbidden to change role for {member} ({member.id}); "
+            f"possibly due to role hierarchy"
+        )
+    except discord.HTTPException as e:
+        log.error(f"Failed to change role for {member} ({member.id}): {e.status} {e.code}")
diff --git a/config-default.yml b/config-default.yml
index 7400cf200..0d3ddc005 100644
--- a/config-default.yml
+++ b/config-default.yml
@@ -264,7 +264,12 @@ guild:
         - *BLACK_FORMATTER
 
     roles:
+        # Self-assignable roles, see the Subscribe cog
+        advent_of_code:                         518565788744024082
         announcements:                          463658397560995840
+        lovefest:                               542431903886606399
+        pyweek_announcements:                   897568414044938310
+
         contributors:                           295488872404484098
         help_cooldown:                          699189276025421825
         muted:              &MUTED_ROLE         277914926603829249