2024-03-18 23:57:13 +00:00
|
|
|
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):
|
2024-03-19 00:20:22 +00:00
|
|
|
|
2024-03-18 23:57:13 +00:00
|
|
|
def __init__(self, bot):
|
|
|
|
self.bot = bot
|
|
|
|
self.quotes_channel_id = CONFIG["quote_a"].get("channel_id")
|
2024-03-19 00:25:45 +00:00
|
|
|
self.names = CONFIG["quote_a"].get("names", {})
|
2024-03-18 23:57:13 +00:00
|
|
|
|
|
|
|
@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(
|
2024-03-19 00:20:22 +00:00
|
|
|
usernames: list[str],
|
|
|
|
counts: list[int],
|
2024-03-19 00:25:45 +00:00
|
|
|
no_other: bool = False
|
2024-03-18 23:57:13 +00:00
|
|
|
) -> 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
|
2024-03-19 00:25:45 +00:00
|
|
|
:param no_other: Disables the "other" grouping
|
2024-03-18 23:57:13 +00:00
|
|
|
:returns: The pie chart image
|
|
|
|
"""
|
2024-03-19 00:14:50 +00:00
|
|
|
|
|
|
|
def pct(v: int):
|
2024-03-19 00:25:45 +00:00
|
|
|
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]
|
2024-03-19 00:20:22 +00:00
|
|
|
|
2024-03-19 00:32:20 +00:00
|
|
|
fig, ax = plt.subplots(figsize=(10, 10))
|
2024-03-18 23:57:13 +00:00
|
|
|
ax.pie(
|
|
|
|
counts,
|
|
|
|
labels=usernames,
|
2024-03-19 00:14:50 +00:00
|
|
|
autopct=pct,
|
2024-03-19 00:25:45 +00:00
|
|
|
startangle=90,
|
2024-03-19 00:30:45 +00:00
|
|
|
radius=1.5,
|
2024-03-18 23:57:13 +00:00
|
|
|
)
|
|
|
|
fio = io.BytesIO()
|
2024-03-19 00:29:21 +00:00
|
|
|
plt.subplots_adjust(left=0.1, bottom=0.1, right=0.9, top=0.9)
|
2024-03-19 00:32:20 +00:00
|
|
|
fig.savefig(fio, format='png')
|
2024-03-18 23:57:13 +00:00
|
|
|
fio.seek(0)
|
2024-03-19 00:32:20 +00:00
|
|
|
return discord.File(fio, filename="pie.png")
|
2024-03-18 23:57:13 +00:00
|
|
|
|
|
|
|
@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
|
|
|
|
)
|
2024-03-19 00:25:45 +00:00
|
|
|
],
|
|
|
|
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
|
|
|
|
)
|
2024-03-18 23:57:13 +00:00
|
|
|
]
|
|
|
|
):
|
|
|
|
"""Checks the quote quota for the quotes channel."""
|
|
|
|
now = discord.utils.utcnow()
|
2024-03-18 23:59:46 +00:00
|
|
|
oldest = now - timedelta(days=days)
|
2024-03-18 23:57:13 +00:00
|
|
|
await ctx.defer()
|
2024-03-18 23:59:46 +00:00
|
|
|
channel = self.quotes_channel or discord.utils.get(ctx.guild.text_channels, name="quotes")
|
2024-03-18 23:57:13 +00:00
|
|
|
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:
|
2024-03-19 00:06:12 +00:00
|
|
|
regex = r".*\s+-\s*@?([\w\s]+)"
|
2024-03-18 23:57:13 +00:00
|
|
|
else:
|
2024-03-19 00:06:12 +00:00
|
|
|
regex = r".+\s*-\s*@?([\w\s]+)"
|
2024-03-18 23:57:13 +00:00
|
|
|
|
2024-03-19 00:06:12 +00:00
|
|
|
if not (m := re.match(regex, str(message.clean_content))):
|
2024-03-18 23:57:13 +00:00
|
|
|
filtered_messages += 1
|
|
|
|
continue
|
|
|
|
name = m.group(1)
|
|
|
|
name = name.strip().title()
|
2024-03-19 00:06:12 +00:00
|
|
|
if name == "Me":
|
2024-03-19 00:25:45 +00:00
|
|
|
name = message.author.name.strip().casefold()
|
|
|
|
if name in self.names:
|
|
|
|
name = self.names[name]
|
|
|
|
else:
|
|
|
|
filtered_messages += 1
|
|
|
|
continue
|
|
|
|
elif name in self.names:
|
|
|
|
name = self.names[name]
|
|
|
|
|
2024-03-18 23:57:13 +00:00
|
|
|
authors.setdefault(name, 0)
|
|
|
|
authors[name] += 1
|
|
|
|
|
|
|
|
file = await asyncio.to_thread(
|
|
|
|
self.generate_pie_chart,
|
|
|
|
list(authors.keys()),
|
2024-03-19 00:25:45 +00:00
|
|
|
list(authors.values()),
|
|
|
|
merge_other
|
2024-03-18 23:57:13 +00:00
|
|
|
)
|
|
|
|
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))
|