Add quota cog

requirements.txt

@@ -18,3 +18,4 @@ humanize~=4.9
 redis~=5.0
 beautifulsoup4~=4.12
 lxml~=5.1
+matplotlib~=3.8

src/cogs/quote_quota.py

@@ -0,0 +1,118 @@
+import asyncio
+import re
+
+import discord
+import io
+import matplotlib.pyplot as plt
+from datetime import timedelta
+from discord.ext import commands
+from typing import Iterable, Annotated
+
+from conf import CONFIG
+
+
+class QuoteQuota(commands.Cog):
+    def __init__(self, bot):
+        self.bot = bot
+        self.quotes_channel_id = CONFIG["quote_a"].get("channel_id")
+
+    @property
+    def quotes_channel(self) -> discord.TextChannel | None:
+        if self.quotes_channel_id:
+            c = self.bot.get_channel(self.quotes_channel_id)
+            if c:
+                return c
+
+    @staticmethod
+    def generate_pie_chart(
+            usernames: Iterable[str],
+            counts: Iterable[int],
+    ) -> discord.File:
+        """
+        Converts the given username and count tuples into a nice pretty pie chart.
+
+        :param usernames: The usernames
+        :param counts: The number of times the username appears in the chat
+        :returns: The pie chart image
+        """
+        fig, ax = plt.subplots()
+        ax.pie(
+            counts,
+            labels=usernames,
+            autopct='%1.1f%%',
+        )
+        ax.legend(loc="upper right")
+        fig.tight_layout()
+        fig.title("Quote Quota")
+        fio = io.BytesIO()
+        fig.savefig(fio, format='jpg')
+        fio.seek(0)
+        return discord.File(fio, filename="pie.jpeg")
+
+    @commands.slash_command()
+    async def quota(
+            self,
+            ctx: discord.ApplicationContext,
+            days: Annotated[
+                int,
+                discord.Option(
+                    int,
+                    name="lookback",
+                    description="How many days to look back on. Defaults to 7.",
+                    default=7,
+                    min_value=1,
+                    max_value=365
+                )
+            ]
+    ):
+        """Checks the quote quota for the quotes channel."""
+        now = discord.utils.utcnow()
+        oldest = now - timedelta(days=7)
+        await ctx.defer()
+        channel = self.quotes_channel
+        if not channel:
+            return await ctx.respond(":x: Cannot find quotes channel.")
+
+        await ctx.respond("Gathering messages, this may take a moment.")
+
+        authors = {}
+        filtered_messages = 0
+        total = 0
+        async for message in channel.history(
+            limit=None,
+            after=oldest,
+            oldest_first=False
+        ):
+            total += 1
+            if not message.content:
+                filtered_messages += 1
+                continue
+            if message.attachments:
+                regex = r".*\s*-\s*(\w+)"
+            else:
+                regex = r".+\s*-\s*(\w+)"
+
+            if not (m := re.match(regex, message.content)):
+                filtered_messages += 1
+                continue
+            name = m.group(1)
+            name = name.strip().title()
+            authors.setdefault(name, 0)
+            authors[name] += 1
+
+        file = await asyncio.to_thread(
+            self.generate_pie_chart,
+            list(authors.keys()),
+            list(authors.values())
+        )
+        return await ctx.edit(
+            content="{:,} messages (out of {:,}) were filtered (didn't follow format?)".format(
+                filtered_messages,
+                total
+            ),
+            file=file
+        )
+
+
+def setup(bot):
+    bot.add_cog(QuoteQuota(bot))

src/cogs/screenshot.py

@@ -168,6 +168,7 @@ class ScreenshotCog(commands.Cog):
         end_save = time.time()
 
         if len(await asyncio.to_thread(file.getvalue)) > 24 * 1024 * 1024:
+            await ctx.edit(content="Compressing screenshot...")
             start_compress = time.time()
             file = await asyncio.to_thread(self.compress_png, file)
             fn = "screenshot.webp"

src/main.py

@@ -133,7 +133,7 @@ bot = Client(
     debug_guilds=CONFIG["jimmy"].get("debug_guilds")
 )
 
-for ext in ("ytdl", "net", "screenshot", "ollama", "ffmeta"):
+for ext in ("ytdl", "net", "screenshot", "ollama", "ffmeta", "quote_quota"):
     try:
         bot.load_extension(f"cogs.{ext}")
     except discord.ExtensionError as e: