Update InfractionScheduler's scheduler to the new API

3 files changed, 12 insertions, 17 deletions

diff --git a/bot/cogs/moderation/management.py b/bot/cogs/moderation/management.py
index c39c7f3bc..e87f3d7a4 100644
--- a/bot/cogs/moderation/management.py
+++ b/bot/cogs/moderation/management.py
@@ -135,11 +135,11 @@ class ModManagement(commands.Cog):
         if 'expires_at' in request_data:
             # A scheduled task should only exist if the old infraction wasn't permanent
             if old_infraction['expires_at']:
-                self.infractions_cog.cancel_task(new_infraction['id'])
+                self.infractions_cog.scheduler.cancel(new_infraction['id'])
 
             # If the infraction was not marked as permanent, schedule a new expiration task
             if request_data['expires_at']:
-                self.infractions_cog.schedule_task(new_infraction['id'], new_infraction)
+                self.infractions_cog.scheduler.schedule(new_infraction['id'], new_infraction)
 
             log_text += f"""
                 Previous expiry: {old_infraction['expires_at'] or "Permanent"}
diff --git a/bot/cogs/moderation/scheduler.py b/bot/cogs/moderation/scheduler.py
index d75a72ddb..601e238c9 100644
--- a/bot/cogs/moderation/scheduler.py
+++ b/bot/cogs/moderation/scheduler.py
@@ -1,4 +1,3 @@
-import asyncio
 import logging
 import textwrap
 import typing as t
@@ -23,13 +22,13 @@ from .utils import UserSnowflake
 log = logging.getLogger(__name__)
 
 
-class InfractionScheduler(Scheduler):
+class InfractionScheduler:
     """Handles the application, pardoning, and expiration of infractions."""
 
     def __init__(self, bot: Bot, supported_infractions: t.Container[str]):
-        super().__init__()
-
         self.bot = bot
+        self.scheduler = Scheduler(self.__class__.__name__)
+
         self.bot.loop.create_task(self.reschedule_infractions(supported_infractions))
 
     @property
@@ -49,7 +48,7 @@ class InfractionScheduler(Scheduler):
         )
         for infraction in infractions:
             if infraction["expires_at"] is not None and infraction["type"] in supported_infractions:
-                self.schedule_task(infraction["id"], infraction)
+                self.schedule_expiration(infraction)
 
     async def reapply_infraction(
         self,
@@ -155,7 +154,7 @@ class InfractionScheduler(Scheduler):
                 await action_coro
                 if expiry:
                     # Schedule the expiration of the infraction.
-                    self.schedule_task(infraction["id"], infraction)
+                    self.schedule_expiration(infraction)
             except discord.HTTPException as e:
                 # Accordingly display that applying the infraction failed.
                 confirm_msg = ":x: failed to apply"
@@ -278,7 +277,7 @@ class InfractionScheduler(Scheduler):
 
                 # Cancel pending expiration task.
                 if infraction["expires_at"] is not None:
-                    self.cancel_task(infraction["id"])
+                    self.scheduler.cancel(infraction["id"])
 
         # Accordingly display whether the user was successfully notified via DM.
         dm_emoji = ""
@@ -415,7 +414,7 @@ class InfractionScheduler(Scheduler):
 
         # Cancel the expiration task.
         if infraction["expires_at"] is not None:
-            self.cancel_task(infraction["id"])
+            self.scheduler.cancel(infraction["id"])
 
         # Send a log message to the mod log.
         if send_log:
@@ -449,7 +448,7 @@ class InfractionScheduler(Scheduler):
         """
         raise NotImplementedError
 
-    async def _scheduled_task(self, infraction: utils.Infraction) -> None:
+    def schedule_expiration(self, infraction: utils.Infraction) -> None:
         """
         Marks an infraction expired after the delay from time of scheduling to time of expiration.
 
@@ -457,8 +456,4 @@ class InfractionScheduler(Scheduler):
         expiration task is cancelled.
         """
         expiry = dateutil.parser.isoparse(infraction["expires_at"]).replace(tzinfo=None)
-        await time.wait_until(expiry)
-
-        # Because deactivate_infraction() explicitly cancels this scheduled task, it is shielded
-        # to avoid prematurely cancelling itself.
-        await asyncio.shield(self.deactivate_infraction(infraction))
+        self.scheduler.schedule_at(expiry, infraction["id"], self.deactivate_infraction(infraction))
diff --git a/bot/cogs/moderation/superstarify.py b/bot/cogs/moderation/superstarify.py
index 45a010f00..867de815a 100644
--- a/bot/cogs/moderation/superstarify.py
+++ b/bot/cogs/moderation/superstarify.py
@@ -146,7 +146,7 @@ class Superstarify(InfractionScheduler, Cog):
         log.debug(f"Changing nickname of {member} to {forced_nick}.")
         self.mod_log.ignore(constants.Event.member_update, member.id)
         await member.edit(nick=forced_nick, reason=reason)
-        self.schedule_task(id_, infraction)
+        self.schedule_expiration(infraction)
 
         # Send a DM to the user to notify them of their new infraction.
         await utils.notify_infraction(