Increase pie size, add name mapping

This commit is contained in:
Nexus 2024-03-19 00:25:45 +00:00
parent 451d9ba77e
commit faf8996a7b
Signed by: nex
GPG key ID: 0FA334385D0B689F

View file

@ -16,6 +16,7 @@ class QuoteQuota(commands.Cog):
def __init__(self, bot): def __init__(self, bot):
self.bot = bot self.bot = bot
self.quotes_channel_id = CONFIG["quote_a"].get("channel_id") self.quotes_channel_id = CONFIG["quote_a"].get("channel_id")
self.names = CONFIG["quote_a"].get("names", {})
@property @property
def quotes_channel(self) -> discord.TextChannel | None: def quotes_channel(self) -> discord.TextChannel | None:
@ -28,36 +29,41 @@ class QuoteQuota(commands.Cog):
def generate_pie_chart( def generate_pie_chart(
usernames: list[str], usernames: list[str],
counts: list[int], counts: list[int],
no_other: bool = False
) -> discord.File: ) -> discord.File:
""" """
Converts the given username and count tuples into a nice pretty pie chart. Converts the given username and count tuples into a nice pretty pie chart.
:param usernames: The usernames :param usernames: The usernames
:param counts: The number of times the username appears in the chat :param counts: The number of times the username appears in the chat
:param no_other: Disables the "other" grouping
:returns: The pie chart image :returns: The pie chart image
""" """
def pct(v: int): def pct(v: int):
return f"{v:.1f}% ({(v / 100) * sum(counts):0f})" return f"{v:.1f}% ({round((v / 100) * sum(counts))})"
other = [] if no_other is False:
# Any authors with less than 5% of the total count will be grouped into "other" other = []
for i, author in enumerate(usernames.copy()): # Any authors with less than 5% of the total count will be grouped into "other"
if (c := counts[i]) / sum(counts) < 0.05: for i, author in enumerate(usernames.copy()):
other.append(c) if (c := counts[i]) / sum(counts) < 0.05:
counts[i] = -1 other.append(c)
usernames.remove(author) counts[i] = -1
if other: usernames.remove(author)
usernames.append("Other") if other:
counts.append(sum(other)) usernames.append("Other")
# And now filter out any -1% counts counts.append(sum(other))
counts = [c for c in counts if c != -1] # And now filter out any -1% counts
counts = [c for c in counts if c != -1]
fig, ax = plt.subplots() fig, ax = plt.subplots()
ax.pie( ax.pie(
counts, counts,
labels=usernames, labels=usernames,
autopct=pct, autopct=pct,
startangle=90,
radius=2
) )
fio = io.BytesIO() fio = io.BytesIO()
fig.savefig(fio, format='jpg') fig.savefig(fio, format='jpg')
@ -78,6 +84,15 @@ class QuoteQuota(commands.Cog):
min_value=1, min_value=1,
max_value=365 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.""" """Checks the quote quota for the quotes channel."""
@ -113,15 +128,23 @@ class QuoteQuota(commands.Cog):
name = m.group(1) name = m.group(1)
name = name.strip().title() name = name.strip().title()
if name == "Me": if name == "Me":
filtered_messages += 1 name = message.author.name.strip().casefold()
continue if name in self.names:
name = self.names[name]
else:
filtered_messages += 1
continue
elif name in self.names:
name = self.names[name]
authors.setdefault(name, 0) authors.setdefault(name, 0)
authors[name] += 1 authors[name] += 1
file = await asyncio.to_thread( file = await asyncio.to_thread(
self.generate_pie_chart, self.generate_pie_chart,
list(authors.keys()), list(authors.keys()),
list(authors.values()) list(authors.values()),
merge_other
) )
return await ctx.edit( return await ctx.edit(
content="{:,} messages (out of {:,}) were filtered (didn't follow format?)".format( content="{:,} messages (out of {:,}) were filtered (didn't follow format?)".format(