diff --git a/requirements.txt b/requirements.txt index da18c0e..0163e13 100644 --- a/requirements.txt +++ b/requirements.txt @@ -18,3 +18,4 @@ humanize~=4.9 redis~=5.0 beautifulsoup4~=4.12 lxml~=5.1 +matplotlib~=3.8 diff --git a/src/cogs/quote_quota.py b/src/cogs/quote_quota.py new file mode 100644 index 0000000..4a3c228 --- /dev/null +++ b/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)) diff --git a/src/cogs/screenshot.py b/src/cogs/screenshot.py index 0f28b24..9b8d518 100644 --- a/src/cogs/screenshot.py +++ b/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" diff --git a/src/main.py b/src/main.py index 82947c0..ef61146 100644 --- a/src/main.py +++ b/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: