Migrate util, closes #113

jarvis/cogs/util.py

@@ -4,93 +4,84 @@ import secrets
 import string
 from io import BytesIO
 
-import discord
-import discord_slash
 import numpy as np
-from discord import File, Guild, Role, User
-from discord.ext import commands
-from discord_slash import SlashContext, cog_ext
-from discord_slash.utils.manage_commands import create_choice, create_option
+from dis_snek import InteractionContext, Scale, Snake, const
+from dis_snek.models.discord.embed import Color, EmbedField
+from dis_snek.models.discord.file import File
+from dis_snek.models.discord.guild import Guild
+from dis_snek.models.discord.role import Role
+from dis_snek.models.discord.user import User
+from dis_snek.models.snek.application_commands import (
+    OptionTypes,
+    SlashCommandChoice,
+    slash_command,
+    slash_option,
+)
+from dis_snek.models.snek.command import cooldown
+from dis_snek.models.snek.cooldowns import Buckets
 from PIL import Image
 
 import jarvis
-from jarvis import jarvis_self
 from jarvis.config import get_config
 from jarvis.data import pigpen
 from jarvis.data.robotcamo import emotes, hk, names
-from jarvis.utils import build_embed, convert_bytesize, get_repo_hash
-from jarvis.utils.field import Field
+from jarvis.utils import build_embed, get_repo_hash
 
 JARVIS_LOGO = Image.open("jarvis_small.png").convert("RGBA")
 
 
-class UtilCog(commands.Cog):
+class UtilCog(Scale):
     """
     Utility functions for J.A.R.V.I.S.
 
     Mostly system utility functions, but may change over time
     """
 
-    def __init__(self, bot: commands.Cog):
+    def __init__(self, bot: Snake):
         self.bot = bot
         self.config = get_config()
 
-    @cog_ext.cog_slash(
-        name="status",
-        description="Retrieve J.A.R.V.I.S. status",
-    )
-    @commands.cooldown(1, 30, commands.BucketType.channel)
-    async def _status(self, ctx: SlashContext) -> None:
+    @slash_command(name="Status", description="Retrieve J.A.R.V.I.S. status")
+    @cooldown(bucket=Buckets.CHANNEL, rate=1, interval=30)
+    async def _status(self, ctx: InteractionContext) -> None:
         title = "J.A.R.V.I.S. Status"
         desc = "All systems online"
-        color = "#98CCDA"
+        color = "#3498db"
         fields = []
-        with jarvis_self.oneshot():
-            fields.append(Field("CPU Usage", jarvis_self.cpu_percent()))
-            fields.append(
-                Field(
-                    "RAM Usage",
-                    convert_bytesize(jarvis_self.memory_info().rss),
-                )
-            )
-            fields.append(Field("PID", jarvis_self.pid))
-            fields.append(Field("discord_slash", discord_slash.__version__))
-            fields.append(Field("discord.py", discord.__version__))
-            fields.append(Field("Version", jarvis.__version__, False))
-            fields.append(Field("Git Hash", get_repo_hash()[:7], False))
-            embed = build_embed(title=title, description=desc, fields=fields, color=color)
-            await ctx.send(embed=embed)
 
-    @cog_ext.cog_slash(
+        fields.append(EmbedField(name="discord.py", value=const.__version__))
+        fields.append(EmbedField(name="Version", value=jarvis.__version__, inline=False))
+        fields.append(EmbedField(name="Git Hash", value=get_repo_hash()[:7], inline=False))
+        embed = build_embed(title=title, description=desc, fields=fields, color=color)
+        await ctx.send(embed=embed)
+
+    @slash_command(
         name="logo",
         description="Get the current logo",
     )
-    @commands.cooldown(1, 30, commands.BucketType.channel)
-    async def _logo(self, ctx: SlashContext) -> None:
+    @cooldown(bucket=Buckets.CHANNEL, rate=1, interval=30)
+    async def _logo(self, ctx: InteractionContext) -> None:
         with BytesIO() as image_bytes:
             JARVIS_LOGO.save(image_bytes, "PNG")
             image_bytes.seek(0)
             logo = File(image_bytes, filename="logo.png")
-
             await ctx.send(file=logo)
 
-    @cog_ext.cog_slash(name="rchk", description="Robot Camo HK416")
-    async def _rchk(self, ctx: SlashContext) -> None:
+    @slash_command(name="rchk", description="Robot Camo HK416")
+    async def _rchk(self, ctx: InteractionContext) -> None:
         await ctx.send(content=hk)
 
-    @cog_ext.cog_slash(
+    @slash_command(
         name="rcauto",
         description="Automates robot camo letters",
-        options=[
-            create_option(
-                name="text",
-                description="Text to camo-ify",
-                option_type=3,
-                required=True,
-            )
-        ],
     )
-    async def _rcauto(self, ctx: SlashContext, text: str) -> None:
+    @slash_option(
+        name="text",
+        description="Text to camo-ify",
+        option_type=OptionTypes.STRING,
+        required=True,
+    )
+    async def _rcauto(self, ctx: InteractionContext, text: str) -> None:
         to_send = ""
         if len(text) == 1 and not re.match(r"^[A-Z0-9-()$@!?^'#. ]$", text.upper()):
             await ctx.send("Please use ASCII characters.", hidden=True)
@@ -109,57 +100,50 @@ class UtilCog(commands.Cog):
         else:
             await ctx.send(to_send)
 
-    @cog_ext.cog_slash(
-        name="avatar",
-        description="Get a user avatar",
-        options=[
-            create_option(
-                name="user",
-                description="User to view avatar of",
-                option_type=6,
-                required=False,
-            )
-        ],
+    @slash_command(name="avatar", description="Get a user avatar")
+    @slash_option(
+        name="user",
+        description="User to view avatar of",
+        option_type=OptionTypes.USER,
+        required=False,
     )
-    @commands.cooldown(1, 5, commands.BucketType.user)
-    async def _avatar(self, ctx: SlashContext, user: User = None) -> None:
+    @cooldown(bucket=Buckets.USER, rate=1, interval=5)
+    async def _avatar(self, ctx: InteractionContext, user: User = None) -> None:
         if not user:
             user = ctx.author
 
         avatar = user.display_avatar
         embed = build_embed(title="Avatar", description="", fields=[], color="#00FFEE")
         embed.set_image(url=avatar)
-        embed.set_author(name=f"{user.name}#{user.discriminator}", icon_url=avatar)
+        embed.set_author(name=f"{user.username}#{user.discriminator}", icon_url=avatar)
         await ctx.send(embed=embed)
 
-    @cog_ext.cog_slash(
+    @slash_command(
         name="roleinfo",
         description="Get role info",
-        options=[
-            create_option(
-                name="role",
-                description="Role to get info of",
-                option_type=8,
-                required=True,
-            )
-        ],
     )
-    async def _roleinfo(self, ctx: SlashContext, role: Role) -> None:
+    @slash_option(
+        name="role",
+        description="Role to get info of",
+        option_type=OptionTypes.ROLE,
+        required=True,
+    )
+    async def _roleinfo(self, ctx: InteractionContext, role: Role) -> None:
         fields = [
-            Field(name="ID", value=role.id),
-            Field(name="Name", value=role.name),
-            Field(name="Color", value=str(role.color)),
-            Field(name="Mention", value=f"`{role.mention}`"),
-            Field(name="Hoisted", value="Yes" if role.hoist else "No"),
-            Field(name="Position", value=str(role.position)),
-            Field(name="Mentionable", value="Yes" if role.mentionable else "No"),
+            EmbedField(name="ID", value=role.id),
+            EmbedField(name="Name", value=role.name),
+            EmbedField(name="Color", value=str(role.color)),
+            EmbedField(name="Mention", value=f"`{role.mention}`"),
+            EmbedField(name="Hoisted", value="Yes" if role.hoist else "No"),
+            EmbedField(name="Position", value=str(role.position)),
+            EmbedField(name="Mentionable", value="Yes" if role.mentionable else "No"),
+            EmbedField(name="Member Count", value=role.members),
         ]
-
         embed = build_embed(
             title="",
             description="",
             fields=fields,
-            color=str(role.color),
+            color=Color.from_hex(role.color),
             timestamp=role.created_at,
         )
         embed.set_footer(text="Role Created")
@@ -181,19 +165,17 @@ class UtilCog(commands.Cog):
 
             await ctx.send(embed=embed, file=color_show)
 
-    @cog_ext.cog_slash(
+    @slash_command(
         name="userinfo",
         description="Get user info",
-        options=[
-            create_option(
-                name="user",
-                description="User to get info of",
-                option_type=6,
-                required=False,
-            )
-        ],
     )
-    async def _userinfo(self, ctx: SlashContext, user: User = None) -> None:
+    @slash_option(
+        name="user",
+        description="User to get info of",
+        option_type=OptionTypes.USER,
+        required=False,
+    )
+    async def _userinfo(self, ctx: InteractionContext, user: User = None) -> None:
         if not user:
             user = ctx.author
         user_roles = user.roles
@@ -201,15 +183,15 @@ class UtilCog(commands.Cog):
             user_roles = sorted(user.roles, key=lambda x: -x.position)
             _ = user_roles.pop(-1)
         fields = [
-            Field(
+            EmbedField(
                 name="Joined",
                 value=user.joined_at.strftime("%a, %b %-d, %Y %-I:%M %p"),
             ),
-            Field(
+            EmbedField(
                 name="Registered",
                 value=user.created_at.strftime("%a, %b %-d, %Y %-I:%M %p"),
             ),
-            Field(
+            EmbedField(
                 name=f"Roles [{len(user_roles)}]",
                 value=" ".join([x.mention for x in user_roles]) if user_roles else "None",
                 inline=False,
@@ -220,24 +202,27 @@ class UtilCog(commands.Cog):
             title="",
             description=user.mention,
             fields=fields,
-            color=str(user_roles[0].color) if user_roles else "#FF0000",
+            color=Color.from_hex(str(user_roles[0].color) if user_roles else "#3498db"),
         )
 
-        embed.set_author(name=f"{user.name}#{user.discriminator}", icon_url=user.display_avatar)
+        embed.set_author(
+            name=f"{user.display_name}#{user.discriminator}", icon_url=user.display_avatar
+        )
         embed.set_thumbnail(url=user.display_avatar)
         embed.set_footer(text=f"ID: {user.id}")
 
         await ctx.send(embed=embed)
 
-    @cog_ext.cog_slash(name="serverinfo", description="Get server info")
-    async def _server_info(self, ctx: SlashContext) -> None:
+    @slash_command(name="serverinfo", description="Get server info")
+    async def _server_info(self, ctx: InteractionContext) -> None:
         guild: Guild = ctx.guild
 
         owner = (
-            f"{guild.owner.name}#{guild.owner.discriminator}" if guild.owner else "||`[redacted]`||"
+            f"{guild.owner.display_name}#{guild.owner.discriminator}"
+            if guild.owner
+            else "||`[redacted]`||"
         )
 
-        region = guild.region
         categories = len(guild.categories)
         text_channels = len(guild.text_channels)
         voice_channels = len(guild.voice_channels)
@@ -246,16 +231,15 @@ class UtilCog(commands.Cog):
         role_list = ", ".join(role.name for role in guild.roles)
 
         fields = [
-            Field(name="Owner", value=owner),
-            Field(name="Region", value=region),
-            Field(name="Channel Categories", value=categories),
-            Field(name="Text Channels", value=text_channels),
-            Field(name="Voice Channels", value=voice_channels),
-            Field(name="Members", value=members),
-            Field(name="Roles", value=roles),
+            EmbedField(name="Owner", value=owner),
+            EmbedField(name="Channel Categories", value=categories),
+            EmbedField(name="Text Channels", value=text_channels),
+            EmbedField(name="Voice Channels", value=voice_channels),
+            EmbedField(name="Members", value=members),
+            EmbedField(name="Roles", value=roles),
         ]
         if len(role_list) < 1024:
-            fields.append(Field(name="Role List", value=role_list, inline=False))
+            fields.append(EmbedField(name="Role List", value=role_list, inline=False))
 
         embed = build_embed(title="", description="", fields=fields, timestamp=guild.created_at)
 
@@ -265,35 +249,32 @@ class UtilCog(commands.Cog):
 
         await ctx.send(embed=embed)
 
-    @cog_ext.cog_subcommand(
-        base="pw",
-        name="gen",
-        base_desc="Password utilites",
+    @slash_command(
+        name="pw",
+        sub_cmd_name="gen",
         description="Generate a secure password",
-        guild_ids=[862402786116763668],
+        scopes=[862402786116763668],
+    )
+    @slash_option(
+        name="length",
+        description="Password length (default 32)",
+        option_type=OptionTypes.INTEGER,
+        required=False,
+    )
+    @slash_option(
+        name="chars",
+        description="Characters to include (default last option)",
+        option_type=OptionTypes.INTEGER,
+        required=False,
         options=[
-            create_option(
-                name="length",
-                description="Password length (default 32)",
-                option_type=4,
-                required=False,
-            ),
-            create_option(
-                name="chars",
-                description="Characters to include (default last option)",
-                option_type=4,
-                required=False,
-                choices=[
-                    create_choice(name="A-Za-z", value=0),
-                    create_choice(name="A-Fa-f0-9", value=1),
-                    create_choice(name="A-Za-z0-9", value=2),
-                    create_choice(name="A-Za-z0-9!@#$%^&*", value=3),
-                ],
-            ),
+            SlashCommandChoice(name="A-Za-z", value=0),
+            SlashCommandChoice(name="A-Fa-f0-9", value=1),
+            SlashCommandChoice(name="A-Za-z0-9", value=2),
+            SlashCommandChoice(name="A-Za-z0-9!@#$%^&*", value=3),
         ],
     )
-    @commands.cooldown(1, 15, type=commands.BucketType.user)
-    async def _pw_gen(self, ctx: SlashContext, length: int = 32, chars: int = 3) -> None:
+    @cooldown(bucket=Buckets.USER, rate=1, interval=15)
+    async def _pw_gen(self, ctx: InteractionContext, length: int = 32, chars: int = 3) -> None:
         if length > 256:
             await ctx.send("Please limit password to 256 characters", hidden=True)
             return
@@ -312,14 +293,11 @@ class UtilCog(commands.Cog):
             hidden=True,
         )
 
-    @cog_ext.cog_slash(
-        name="pigpen",
-        description="Encode a string into pigpen",
-        options=[
-            create_option(name="text", description="Text to encode", option_type=3, required=True)
-        ],
+    @slash_command(name="pigpen", description="Encode a string into pigpen")
+    @slash_option(
+        name="text", description="Text to encode", option_type=OptionTypes.STRING, required=True
     )
-    async def _pigpen(self, ctx: SlashContext, text: str) -> None:
+    async def _pigpen(self, ctx: InteractionContext, text: str) -> None:
         outp = "`"
         for c in text:
             c = c.lower()
@@ -334,6 +312,6 @@ class UtilCog(commands.Cog):
         await ctx.send(outp[:2000])
 
 
-def setup(bot: commands.Bot) -> None:
+def setup(bot: Snake) -> None:
     """Add UtilCog to J.A.R.V.I.S."""
     bot.add_cog(UtilCog(bot))