Migrate verify, closes #105

jarvis/cogs/verify.py

@@ -1,10 +1,12 @@
 """J.A.R.V.I.S. Verify Cog."""
+import asyncio
 from random import randint
 
-from discord.ext import commands
-from discord_slash import ComponentContext, SlashContext, cog_ext
-from discord_slash.model import ButtonStyle
-from discord_slash.utils import manage_components
+from dis_snek import InteractionContext, Scale, Snake
+from dis_snek.models.application_commands import slash_command
+from dis_snek.models.discord.components import Button, ButtonStyles, spread_to_rows
+from dis_snek.models.snek.command import cooldown
+from dis_snek.models.snek.cooldowns import Buckets
 
 from jarvis.db.models import Setting
 
@@ -16,30 +18,27 @@ def create_layout() -> list:
     for i in range(3):
         label = "YES" if i == yes else "NO"
         id = f"no_{i}" if not i == yes else "yes"
-        color = ButtonStyle.green if i == yes else ButtonStyle.red
+        color = ButtonStyles.GREEN if i == yes else ButtonStyles.RED
         buttons.append(
-            manage_components.create_button(
+            Button(
                 style=color,
                 label=label,
                 custom_id=f"verify_button||{id}",
             )
         )
-    action_row = manage_components.spread_to_rows(*buttons, max_in_row=3)
+    action_row = spread_to_rows(*buttons, max_in_row=3)
     return action_row
 
 
-class VerifyCog(commands.Cog):
+class VerifyCog(Scale):
     """J.A.R.V.I.S. Verify Cog."""
 
-    def __init__(self, bot: commands.Bot):
+    def __init__(self, bot: Snake):
         self.bot = bot
 
-    @cog_ext.cog_slash(
-        name="verify",
-        description="Verify that you've read the rules",
-    )
-    @commands.cooldown(1, 15, commands.BucketType.user)
-    async def _verify(self, ctx: SlashContext) -> None:
+    @slash_command(name="verify", description="Verify that you've read the rules")
+    @cooldown(bucket=Buckets.USER, rate=1, interval=15)
+    async def _verify(self, ctx: InteractionContext) -> None:
         await ctx.defer()
         role = Setting.objects(guild=ctx.guild.id, setting="verified").first()
         if not role:
@@ -53,40 +52,41 @@ class VerifyCog(commands.Cog):
             content=f"{ctx.author.mention}, please press the button that says `YES`.",
             components=components,
         )
-        await message.delete(delay=15)
 
-    @cog_ext.cog_component(components=create_layout())
-    async def _process(self, ctx: ComponentContext) -> None:
-        await ctx.defer(edit_origin=True)
         try:
-            if ctx.author.id != ctx.origin_message.mentions[0].id:
-                return
-        except Exception:
-            return
-        correct = ctx.custom_id.split("||")[-1] == "yes"
-        if correct:
-            components = ctx.origin_message.components
-            for c in components:
-                for c2 in c["components"]:
-                    c2["disabled"] = True
-            setting = Setting.objects(guild=ctx.guild.id, setting="verified").first()
-            role = await ctx.guild.get_role(setting.value)
-            await ctx.author.add_roles(role, reason="Verification passed")
-            setting = Setting.objects(guild=ctx.guild.id, setting="unverified").first()
-            if setting:
+            context = await self.bot.wait_for_component(
+                messages=message, check=lambda x: ctx.author.id == x.author.id, timeout=30
+            )
+
+            correct = context.context.custom_id.split("||")[-1] == "yes"
+            if correct:
+                for row in components:
+                    for component in row["components"]:
+                        component["disabled"] = True
+                setting = Setting.objects(guild=ctx.guild.id, setting="verified").first()
                 role = await ctx.guild.get_role(setting.value)
-                await ctx.author.remove_roles(role, reason="Verification passed")
-            await ctx.edit_origin(
-                content=f"Welcome, {ctx.author.mention}. Please enjoy your stay.",
-                components=manage_components.spread_to_rows(*components, max_in_row=5),
-            )
-            await ctx.origin_message.delete(delay=5)
-        else:
-            await ctx.edit_origin(
-                content=f"{ctx.author.mention}, incorrect. Please press the button that says `YES`",
-            )
+                await ctx.author.add_roles(role, reason="Verification passed")
+                setting = Setting.objects(guild=ctx.guild.id, setting="unverified").first()
+                if setting:
+                    role = await ctx.guild.get_role(setting.value)
+                    await ctx.author.remove_roles(role, reason="Verification passed")
+
+                await context.context.edit_origin(
+                    content=f"Welcome, {ctx.author.mention}. Please enjoy your stay.",
+                    components=components,
+                )
+                await context.context.message.delete(delay=5)
+            else:
+                await context.context.edit_origin(
+                    content=(
+                        f"{ctx.author.mention}, incorrect. "
+                        "Please press the button that says `YES`"
+                    )
+                )
+        except asyncio.TimeoutError:
+            await message.delete(delay=30)
 
 
-def setup(bot: commands.Bot) -> None:
+def setup(bot: Snake) -> None:
     """Add VerifyCog to J.A.R.V.I.S."""
     bot.add_cog(VerifyCog(bot))