Merge remote-tracking branch 'origin/master'

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,188 @@
+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")
+        self.names = CONFIG["quote_a"].get("names", {})
+
+    @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: list[str],
+            counts: list[int],
+            no_other: bool = False
+    ) -> 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
+        :param no_other: Disables the "other" grouping
+        :returns: The pie chart image
+        """
+
+        def pct(v: int):
+            return f"{v:.1f}% ({round((v / 100) * sum(counts))})"
+
+        if no_other is False:
+            other = []
+            # Any authors with less than 5% of the total count will be grouped into "other"
+            for i, author in enumerate(usernames.copy()):
+                if (c := counts[i]) / sum(counts) < 0.05:
+                    other.append(c)
+                    counts[i] = -1
+                    usernames.remove(author)
+            if other:
+                usernames.append("Other")
+                counts.append(sum(other))
+            # And now filter out any -1% counts
+            counts = [c for c in counts if c != -1]
+
+        mapping = {}
+        for i, author in enumerate(usernames):
+            mapping[author] = counts[i]
+
+        # Sort the authors by count
+        new_mapping = {}
+        for author, count in sorted(mapping.items(), key=lambda x: x[1], reverse=True):
+            new_mapping[author] = count
+
+        usernames = list(new_mapping.keys())
+        counts = list(new_mapping.values())
+
+        fig, ax = plt.subplots(figsize=(7, 7))
+        ax.pie(
+            counts,
+            labels=usernames,
+            autopct=pct,
+            startangle=90,
+            radius=1.2,
+        )
+        fig.subplots_adjust(left=0.1, bottom=0.1, right=0.9, top=0.9, wspace=0.3, hspace=0.4)
+        fio = io.BytesIO()
+        fig.savefig(fio, format='png')
+        fio.seek(0)
+        return discord.File(fio, filename="pie.png")
+
+    @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
+                )
+            ],
+            merge_other: Annotated[
+                bool,
+                discord.Option(
+                    bool,
+                    name="merge_other",
+                    description="Whether to merge authors with less than 5% of the total count into 'Other'.",
+                    default=True
+                )
+            ]
+    ):
+        """Checks the quote quota for the quotes channel."""
+        now = discord.utils.utcnow()
+        oldest = now - timedelta(days=days)
+        await ctx.defer()
+        channel = self.quotes_channel or discord.utils.get(ctx.guild.text_channels, name="quotes")
+        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\s]+)"
+            else:
+                regex = r".+\s+-\s*@?([\w\s]+)"
+
+            if not (m := re.match(regex, str(message.clean_content))):
+                filtered_messages += 1
+                continue
+            name = m.group(1)
+            name = name.strip().casefold()
+            if name == "me":
+                name = message.author.name.strip().casefold()
+                if name in self.names:
+                    name = self.names[name].title()
+                else:
+                    filtered_messages += 1
+                    continue
+            elif name in self.names:
+                name = self.names[name].title()
+            elif name.isdigit():
+                filtered_messages += 1
+                continue
+
+            name = name.title()
+            authors.setdefault(name, 0)
+            authors[name] += 1
+
+        if not authors:
+            if total:
+                return await ctx.edit(
+                    content="No valid messages found in the last {!s} days. "
+                            "Make sure quotes are formatted properly ending with ` - AuthorName`"
+                            " (e.g. `\"This is my quote\" - Jimmy`)".format(days)
+                )
+            else:
+                return await ctx.edit(
+                    content="No messages found in the last {!s} days.".format(days)
+                )
+
+        file = await asyncio.to_thread(
+            self.generate_pie_chart,
+            list(authors.keys()),
+            list(authors.values()),
+            merge_other
+        )
+        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

@@ -5,6 +5,7 @@ import logging
 import os
 import tempfile
 import time
+import copy
 from urllib.parse import urlparse
 
 import discord
@@ -15,6 +16,8 @@ from selenium import webdriver
 from selenium.webdriver.chrome.options import Options as ChromeOptions
 from selenium.webdriver.chrome.service import Service as ChromeService
 
+from conf import CONFIG
+
 
 class ScreenshotCog(commands.Cog):
     def __init__(self, bot: commands.Bot):
@@ -76,7 +79,8 @@ class ScreenshotCog(commands.Cog):
             load_timeout: int = 10,
             render_timeout: int = None,
             eager: bool = None,
-            resolution: str = "1920x1080"
+            resolution: str = "1920x1080",
+            use_proxy: bool = False
     ):
         """Screenshots a webpage."""
         await ctx.defer()
@@ -104,11 +108,14 @@ class ScreenshotCog(commands.Cog):
 
         start_init = time.time()
         try:
+            options = copy.copy(self.chrome_options)
+            if use_proxy and (server := CONFIG["screenshot"].get("proxy")):
+                options.add_argument("--proxy-server=" + server)
             service = await asyncio.to_thread(ChromeService)
             driver: webdriver.Chrome = await asyncio.to_thread(
                 webdriver.Chrome,
                 service=service,
-                options=self.chrome_options
+                options=options
             )
             driver.set_page_load_timeout(load_timeout)
             if resolution:
@@ -161,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/conf.py

@@ -28,6 +28,8 @@ try:
     CONFIG.setdefault("jimmy", {})
     CONFIG.setdefault("ollama", {})
     CONFIG.setdefault("rss", {"meta": {"channel": None}})
+    CONFIG.setdefault("screenshot", {})
+    CONFIG.setdefault("quote_a", {"channel": None})
     CONFIG.setdefault(
         "server",
         {

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: