Allow certain roles to run AoC commands in November (#171)

* Allow certain roles to run AoC commands in November

* Fix role condition for listeners

Co-authored-by: lemonyte <49930425+lemonyte@users.noreply.github.com>

* Use MissingAnyRole instead of generic CheckFailure

---------

Co-authored-by: lemonyte <49930425+lemonyte@users.noreply.github.com>

Author: mbaruh

bot/constants.py

@@ -154,6 +154,7 @@ class _Roles(EnvConfig, env_prefix="ROLE_"):
     admins: int = 267628507062992896
     advent_of_code: int = 518565788744024082
     code_jam_event_team: int = 787816728474288181
+    core_devs: int = 587606783669829632
     events_lead: int = 778361735739998228
     event_runner: int = 940911658799333408
     summer_aoc: int = 988801794668908655

bot/exts/advent_of_code/_cog.py

@@ -37,6 +37,8 @@
 
 AOC_REDIRECT = (Channels.advent_of_code_commands, Channels.sir_lancebot_playground, Channels.bot_commands)
 
+EARLY_ACCESS_ROLES = (Roles.admins, Roles.events_lead, Roles.core_devs)
+
 
 class AdventOfCode(commands.Cog):
     """Advent of Code festivities! Ho Ho Ho!"""
@@ -386,6 +388,7 @@ async def aoc_unlink_account(self, ctx: commands.Context) -> None:
             await ctx.reply("You don't have an Advent of Code account linked.")
 
     @in_month(Month.DECEMBER, Month.JANUARY, Month.FEBRUARY)
+    @in_month(Month.NOVEMBER, roles=EARLY_ACCESS_ROLES)
     @adventofcode_group.command(
         name="dayandstar",
         aliases=("daynstar", "daystar"),
@@ -424,6 +427,7 @@ async def aoc_day_and_star_leaderboard(
         await message.edit(view=None)
 
     @in_month(Month.DECEMBER, Month.JANUARY, Month.FEBRUARY)
+    @in_month(Month.NOVEMBER, roles=EARLY_ACCESS_ROLES)
     @adventofcode_group.command(
         name="leaderboard",
         aliases=("board", "lb"),
@@ -505,6 +509,7 @@ async def aoc_global_leaderboard(self, ctx: commands.Context) -> None:
         await ctx.send(embed=self.cached_no_global)
 
     @in_month(Month.DECEMBER, Month.JANUARY, Month.FEBRUARY)
+    @in_month(Month.NOVEMBER, roles=EARLY_ACCESS_ROLES)
     @adventofcode_group.command(
         name="stats",
         aliases=("dailystats", "ds"),

bot/utils/decorators.py

@@ -4,7 +4,7 @@
 from collections.abc import Callable, Container
 
 from discord.ext import commands
-from discord.ext.commands import Command, Context
+from discord.ext.commands import Command, Context, errors
 from pydis_core.utils import logging
 
 from bot.constants import Channels
@@ -70,28 +70,66 @@ async def guarded_listener(*args, **kwargs) -> None:
     return decorator
 
 
-def in_month_command(*allowed_months: Month) -> Callable:
+def in_month_command(*allowed_months: Month, roles: tuple[int, ...] = ()) -> Callable:
     """
-    Check whether the command was invoked in one of `enabled_months`.
+    Check whether the command was invoked in one of `allowed_months`.
+
+    The check can be limited to certain roles.
+    To enable the supplied months for everyone, don't set a value for `roles`.
+
+    If a command is decorated several times with this, it only needs to pass one of the checks.
 
     Uses the current UTC month at the time of running the predicate.
     """
     async def predicate(ctx: Context) -> bool:
-        current_month = resolve_current_month()
-        can_run = current_month in allowed_months
-
-        log.debug(
-            f"Command '{ctx.command}' is locked to months {human_months(allowed_months)}. "
-            f"Invoking it in month {current_month!s} is {'allowed' if can_run else 'disallowed'}."
-        )
-        if can_run:
+        command = ctx.command
+        if "month_checks" not in command.extras:
+            log.debug(f"No month checks found for command {command}.")
             return True
-        raise InMonthCheckFailure(f"Command can only be used in {human_months(allowed_months)}")
 
-    return commands.check(predicate)
+        everyone_error = None
+        privileged_user = False
+        allowed_months_for_user = set()
+        current_month = resolve_current_month()
+        user_roles = set(r.id for r in ctx.author.roles) if ctx.author.roles else set()
+
+        for checked_roles, checked_months in command.extras["month_checks"].items():
+            if checked_roles:
+                if not set(checked_roles) & user_roles:
+                    log.debug(f"Month check for roles {checked_roles} doesn't apply to {ctx.author}.")
+                    continue
+
+            if current_month in checked_months:
+                log.debug(f"Month check for roles {checked_roles} passed for {ctx.author}.")
+                return True
+
+            log.debug(f"Month check for roles {checked_roles} didn't pass for {ctx.author}.")
+            if not checked_roles:
+                everyone_error = InMonthCheckFailure(f"Command can only be used in {human_months(checked_months)}")
+            else:
+                privileged_user = True
+            allowed_months_for_user |= set(checked_months)
+
+        if privileged_user:
+            allowed_months_for_user = sorted(allowed_months_for_user)
+            raise InMonthCheckFailure(f"You can run this command only in {human_months(allowed_months_for_user)}")
+        if everyone_error:
+            raise everyone_error
+        # There's no general access to this command, and the user has no relevant roles.
+        raise errors.MissingAnyRole(list({r for rs in command.extras["month_checks"] for r in rs}))
+
+    def decorator(func: Command) -> Command:
+        if "month_checks" in func.extras:
+            func.extras["month_checks"][roles] = allowed_months
+            return func
+
+        func.extras["month_checks"] = {roles: allowed_months}
+        return commands.check(predicate)(func)
+
+    return decorator
 
 
-def in_month(*allowed_months: Month) -> Callable:
+def in_month(*allowed_months: Month, roles: tuple[int, ...] = ()) -> Callable:
     """
     Universal decorator for season-locking commands and listeners alike.
 
@@ -113,10 +151,12 @@ def decorator(callable_: Callable) -> Callable:
         # Functions decorated as commands are turned into instances of `Command`
         if isinstance(callable_, Command):
             log.debug(f"Command {callable_.qualified_name} will be locked to {human_months(allowed_months)}")
-            actual_deco = in_month_command(*allowed_months)
+            actual_deco = in_month_command(*allowed_months, roles=roles)
 
         # D.py will assign this attribute when `callable_` is registered as a listener
         elif hasattr(callable_, "__cog_listener__"):
+            if roles:
+                raise ValueError("Role restrictions are not available for listeners.")
             log.debug(f"Listener {callable_.__qualname__} will be locked to {human_months(allowed_months)}")
             actual_deco = in_month_listener(*allowed_months)