Migrate image cog

jarvis/cogs/image.py

@@ -5,13 +5,16 @@ from io import BytesIO
 import aiohttp
 import cv2
 import numpy as np
-from dis_snek import MessageContext, Scale, Snake, message_command
+from dis_snek import InteractionContext, Scale, Snake
 from dis_snek.models.discord.embed import EmbedField
 from dis_snek.models.discord.file import File
-from dis_snek.models.snek.command import cooldown
+from dis_snek.models.discord.message import Attachment
-from dis_snek.models.snek.cooldowns import Buckets
+from dis_snek.models.snek.application_commands import (
-
+    OptionTypes,
-from jarvis.utils import build_embed, convert_bytesize, unconvert_bytesize
+    slash_command,
+    slash_option,
+)
+from jarvis_core.util import build_embed, convert_bytesize, unconvert_bytesize
 
 MIN_ACCURACY = 0.80
 
@@ -26,39 +29,65 @@ class ImageCog(Scale):
     def __init__(self, bot: Snake):
         self.bot = bot
         self._session = aiohttp.ClientSession()
-        self.tgt_match = re.compile(r"([0-9]*\.?[0-9]*?) ?([KMGTP]?B)", re.IGNORECASE)
+        self.tgt_match = re.compile(r"([0-9]*\.?[0-9]*?) ?([KMGTP]?B?)", re.IGNORECASE)
 
     def __del__(self):
         self._session.close()
 
-    async def _resize(self, ctx: MessageContext, target: str, url: str = None) -> None:
+    @slash_command(name="resize", description="Resize an image")
-        if not target:
+    @slash_option(
-            await ctx.send("Missing target size, i.e. 200KB.")
+        name="target",
+        description="Target size, i.e. 200KB",
+        opt_type=OptionTypes.STRING,
+        required=True,
+    )
+    @slash_option(
+        name="attachment",
+        description="Image to resize",
+        opt_type=OptionTypes.ATTACHMENT,
+        required=False,
+    )
+    @slash_option(
+        name="url",
+        description="URL to download and resize",
+        opt_type=OptionTypes.STRING,
+        required=False,
+    )
+    async def _resize(
+        self, ctx: InteractionContext, target: str, attachment: Attachment = None, url: str = None
+    ) -> None:
+        if not attachment and not url:
+            await ctx.send("A URL or attachment is required", ephemeral=True)
+            return
+
+        if attachment and not attachment.content_type.startswith("image"):
+            await ctx.send("Attachment must be an image", ephemeral=True)
             return
 
         tgt = self.tgt_match.match(target)
         if not tgt:
-            await ctx.send(f"Invalid target format ({target}). Expected format like 200KB")
+            await ctx.send(
+                f"Invalid target format ({target}). Expected format like 200KB", ephemeral=True
+            )
             return
 
         tgt_size = unconvert_bytesize(float(tgt.groups()[0]), tgt.groups()[1])
         if tgt_size > unconvert_bytesize(8, "MB"):
-            await ctx.send("Target too large to send. Please make target < 8MB")
+            await ctx.send("Target too large to send. Please make target < 8MB", ephemeral=True)
             return
-        file = None
-        filename = None
-        if ctx.message.attachments is not None and len(ctx.message.attachments) > 0:
-            file = await ctx.message.attachments[0].read()
-            filename = ctx.message.attachments[0].filename
-        elif url is not None:
-            async with self._session.get(url) as resp:
-                if resp.status == 200:
-                    file = await resp.read()
-            filename = url.split("/")[-1]
-        else:
-            ctx.send("Missing file as either attachment or URL.")
 
-        size = len(file)
+        if attachment:
+            url = attachment.url
+            filename = attachment.filename
+        else:
+            filename = url.split("/")[-1]
+
+        data = None
+        async with self._session.get(url) as resp:
+            if resp.status == 200:
+                data = await resp.read()
+
+        size = len(data)
         if size <= tgt_size:
             await ctx.send("Image already meets target.")
             return
@@ -67,43 +96,38 @@ class ImageCog(Scale):
 
         accuracy = 0.0
 
-        while len(file) > tgt_size or (len(file) <= tgt_size and accuracy < MIN_ACCURACY):
+        while len(data) > tgt_size or (len(data) <= tgt_size and accuracy < MIN_ACCURACY):
-            old_file = file
+            old_file = data
 
-            buffer = np.frombuffer(file, dtype=np.uint8)
+            buffer = np.frombuffer(data, dtype=np.uint8)
             img = cv2.imdecode(buffer, flags=-1)
 
             width = int(img.shape[1] * ratio)
             height = int(img.shape[0] * ratio)
 
             new_img = cv2.resize(img, (width, height))
-            file = cv2.imencode(".png", new_img)[1].tobytes()
+            data = cv2.imencode(".png", new_img)[1].tobytes()
-            accuracy = (len(file) / tgt_size) * 100
+            accuracy = (len(data) / tgt_size) * 100
             if accuracy <= 0.50:
-                file = old_file
+                data = old_file
                 ratio += 0.1
             else:
-                ratio = max(tgt_size / len(file) - 0.02, 0.65)
+                ratio = max(tgt_size / len(data) - 0.02, 0.65)
 
-        bufio = BytesIO(file)
+        bufio = BytesIO(data)
-        accuracy = (len(file) / tgt_size) * 100
+        accuracy = (len(data) / tgt_size) * 100
         fields = [
             EmbedField("Original Size", convert_bytesize(size), False),
-            EmbedField("New Size", convert_bytesize(len(file)), False),
+            EmbedField("New Size", convert_bytesize(len(data)), False),
             EmbedField("Accuracy", f"{accuracy:.02f}%", False),
         ]
         embed = build_embed(title=filename, description="", fields=fields)
         embed.set_image(url="attachment://resized.png")
         await ctx.send(
             embed=embed,
-            file=File(file=bufio, filename="resized.png"),
+            file=File(file=bufio, file_name="resized.png"),
         )
 
-    @message_command(name="resize")
-    @cooldown(bucket=Buckets.USER, rate=1, interval=60)
-    async def _resize_pref(self, ctx: MessageContext, target: str, url: str = None) -> None:
-        await self._resize(ctx, target, url)
-
 
 def setup(bot: Snake) -> None:
     """Add ImageCog to J.A.R.V.I.S."""