366 lines
14 KiB
Python
366 lines
14 KiB
Python
# This cog is influenced heavily by cacobot's stats module:
|
|
# https://github.com/Orangestar12/cacobot/blob/master/cacobot/stats.py
|
|
# Big thanks to Redjumpman for changing the beta version from
|
|
# Imagemagick/cairosvg to matplotlib.
|
|
# Thanks to violetnyte for suggesting this cog.
|
|
|
|
import asyncio
|
|
import discord
|
|
import heapq
|
|
from io import BytesIO
|
|
from typing import List, Optional, Tuple, Union
|
|
|
|
from redbot.core import checks, commands, Config
|
|
|
|
import matplotlib
|
|
matplotlib.use("agg")
|
|
import matplotlib.pyplot as plt
|
|
plt.switch_backend("agg")
|
|
|
|
|
|
class Chatchart(commands.Cog):
|
|
"""Show activity."""
|
|
|
|
async def red_delete_data_for_user(self, **kwargs):
|
|
""" Nothing to delete """
|
|
return
|
|
|
|
def __init__(self, bot):
|
|
self.bot = bot
|
|
self.config = Config.get_conf(self, 2766691001, force_registration=True)
|
|
|
|
default_guild = {"channel_deny": []}
|
|
default_global = {"limit": 0}
|
|
|
|
self.config.register_guild(**default_guild)
|
|
self.config.register_global(**default_global)
|
|
|
|
@staticmethod
|
|
def calculate_member_perc(history: List[discord.Message]) -> dict:
|
|
"""Calculate the member count from the message history"""
|
|
msg_data = {"total_count": 0, "users": {}}
|
|
for msg in history:
|
|
# Name formatting
|
|
if len(msg.author.display_name) >= 20:
|
|
short_name = "{}...".format(msg.author.display_name[:20]).replace("$", "\\$")
|
|
else:
|
|
short_name = msg.author.display_name.replace("$", "\\$").replace("_", "\\_ ").replace("*", "\\*")
|
|
whole_name = "{}#{}".format(short_name, msg.author.discriminator)
|
|
if msg.author.bot:
|
|
pass
|
|
elif whole_name in msg_data["users"]:
|
|
msg_data["users"][whole_name]["msgcount"] += 1
|
|
msg_data["total_count"] += 1
|
|
else:
|
|
msg_data["users"][whole_name] = {}
|
|
msg_data["users"][whole_name]["msgcount"] = 1
|
|
msg_data["total_count"] += 1
|
|
return msg_data
|
|
|
|
@staticmethod
|
|
def calculate_top(msg_data: dict) -> Tuple[list, int]:
|
|
"""Calculate the top 20 from the message data package"""
|
|
for usr in msg_data["users"]:
|
|
pd = float(msg_data["users"][usr]["msgcount"]) / float(msg_data["total_count"])
|
|
msg_data["users"][usr]["percent"] = pd * 100
|
|
top_twenty = heapq.nlargest(
|
|
20,
|
|
[
|
|
(x, msg_data["users"][x][y])
|
|
for x in msg_data["users"]
|
|
for y in msg_data["users"][x]
|
|
if (y == "percent" and msg_data["users"][x][y] > 0)
|
|
],
|
|
key=lambda x: x[1],
|
|
)
|
|
others = 100 - sum(x[1] for x in top_twenty)
|
|
return top_twenty, others
|
|
|
|
@staticmethod
|
|
async def create_chart(top, others, channel_or_guild: Union[discord.Guild, discord.TextChannel]):
|
|
plt.clf()
|
|
sizes = [x[1] for x in top]
|
|
labels = ["{} {:g}%".format(x[0], round(x[1], 1)) for x in top]
|
|
if len(top) >= 20:
|
|
sizes = sizes + [others]
|
|
labels = labels + ["Others {:g}%".format(round(others, 1))]
|
|
if len(channel_or_guild.name) >= 19:
|
|
if isinstance(channel_or_guild, discord.Guild):
|
|
channel_or_guild_name = "{}...".format(channel_or_guild.name[:19])
|
|
else:
|
|
channel_or_guild_name = "#{}...".format(channel_or_guild.name[:19])
|
|
else:
|
|
channel_or_guild_name = channel_or_guild.name
|
|
title = plt.title("Stats in {}".format(channel_or_guild_name), color="white")
|
|
title.set_va("top")
|
|
title.set_ha("center")
|
|
plt.gca().axis("equal")
|
|
colors = [
|
|
"r",
|
|
"darkorange",
|
|
"gold",
|
|
"y",
|
|
"olivedrab",
|
|
"green",
|
|
"darkcyan",
|
|
"mediumblue",
|
|
"darkblue",
|
|
"blueviolet",
|
|
"indigo",
|
|
"orchid",
|
|
"mediumvioletred",
|
|
"crimson",
|
|
"chocolate",
|
|
"yellow",
|
|
"limegreen",
|
|
"forestgreen",
|
|
"dodgerblue",
|
|
"slateblue",
|
|
"gray",
|
|
]
|
|
pie = plt.pie(sizes, colors=colors, startangle=0)
|
|
plt.legend(
|
|
pie[0],
|
|
labels,
|
|
bbox_to_anchor=(0.7, 0.5),
|
|
loc="center",
|
|
fontsize=10,
|
|
bbox_transform=plt.gcf().transFigure,
|
|
facecolor="#ffffff",
|
|
)
|
|
plt.subplots_adjust(left=0.0, bottom=0.1, right=0.45)
|
|
image_object = BytesIO()
|
|
plt.savefig(image_object, format="PNG", facecolor="#36393E")
|
|
image_object.seek(0)
|
|
return image_object
|
|
|
|
async def fetch_channel_history(
|
|
self,
|
|
channel: discord.TextChannel,
|
|
animation_message: discord.Message,
|
|
messages: int
|
|
) -> List[discord.Message]:
|
|
"""Fetch the history of a channel while displaying an status message with it"""
|
|
animation_message_deleted = False
|
|
history = []
|
|
history_counter = 0
|
|
async for msg in channel.history(limit=messages):
|
|
history.append(msg)
|
|
history_counter += 1
|
|
await asyncio.sleep(0.005)
|
|
if history_counter % 250 == 0:
|
|
new_embed = discord.Embed(
|
|
title=f"Fetching messages from #{channel.name}",
|
|
description=f"This might take a while...\n{history_counter}/{messages} messages gathered",
|
|
colour=await self.bot.get_embed_colour(location=channel),
|
|
)
|
|
if channel.permissions_for(channel.guild.me).send_messages:
|
|
await channel.typing()
|
|
if animation_message_deleted is False:
|
|
try:
|
|
await animation_message.edit(embed=new_embed)
|
|
except discord.NotFound:
|
|
animation_message_deleted = True
|
|
return history
|
|
|
|
@commands.guild_only()
|
|
@commands.command()
|
|
@commands.cooldown(1, 10, commands.BucketType.guild)
|
|
@commands.max_concurrency(1, commands.BucketType.guild)
|
|
@commands.bot_has_permissions(attach_files=True)
|
|
async def chatchart(self, ctx, channel: Optional[discord.TextChannel] = None, messages:int = 5000):
|
|
"""
|
|
Generates a pie chart, representing the last 5000 messages in the specified channel.
|
|
"""
|
|
if channel is None:
|
|
channel = ctx.channel
|
|
|
|
# --- Early terminations
|
|
if channel.permissions_for(ctx.message.author).read_messages is False:
|
|
return await ctx.send("You're not allowed to access that channel.")
|
|
if channel.permissions_for(ctx.guild.me).read_messages is False:
|
|
return await ctx.send("I cannot read the history of that channel.")
|
|
blacklisted_channels = await self.config.guild(ctx.guild).channel_deny()
|
|
if channel.id in blacklisted_channels:
|
|
return await ctx.send(f"I am not allowed to create a chatchart of {channel.mention}.")
|
|
if messages < 5:
|
|
return await ctx.send("Don't be silly.")
|
|
|
|
message_limit = await self.config.limit()
|
|
if (message_limit != 0) and (messages > message_limit):
|
|
messages = message_limit
|
|
|
|
embed = discord.Embed(
|
|
title=f"Fetching messages from #{channel.name}",
|
|
description="This might take a while...",
|
|
colour=await self.bot.get_embed_colour(location=channel)
|
|
)
|
|
loading_message = await ctx.send(embed=embed)
|
|
try:
|
|
history = await self.fetch_channel_history(channel, loading_message, messages)
|
|
except discord.errors.Forbidden:
|
|
try:
|
|
await loading_message.delete()
|
|
except discord.NotFound:
|
|
pass
|
|
return await ctx.send("No permissions to read that channel.")
|
|
|
|
msg_data = self.calculate_member_perc(history)
|
|
# If no members are found.
|
|
if len(msg_data["users"]) == 0:
|
|
try:
|
|
await loading_message.delete()
|
|
except discord.NotFound:
|
|
pass
|
|
return await ctx.send(f"Only bots have sent messages in {channel.mention} or I can't read message history.")
|
|
|
|
top_twenty, others = self.calculate_top(msg_data)
|
|
chart = await self.create_chart(top_twenty, others, channel)
|
|
|
|
try:
|
|
await loading_message.delete()
|
|
except discord.NotFound:
|
|
pass
|
|
await ctx.send(file=discord.File(chart, "chart.png"))
|
|
|
|
@checks.mod_or_permissions(manage_guild=True)
|
|
@commands.guild_only()
|
|
@commands.command(aliases=["guildchart"])
|
|
@commands.cooldown(1, 30, commands.BucketType.guild)
|
|
@commands.max_concurrency(1, commands.BucketType.guild)
|
|
@commands.bot_has_permissions(attach_files=True)
|
|
async def serverchart(self, ctx: commands.Context, messages: int = 1000):
|
|
"""
|
|
Generates a pie chart, representing the last 1000 messages from every allowed channel in the server.
|
|
|
|
As example:
|
|
For each channel that the bot is allowed to scan. It will take the last 1000 messages from each channel.
|
|
And proceed to build a chart out of that.
|
|
"""
|
|
if messages < 5:
|
|
return await ctx.send("Don't be silly.")
|
|
channel_list = []
|
|
blacklisted_channels = await self.config.guild(ctx.guild).channel_deny()
|
|
for channel in ctx.guild.text_channels:
|
|
channel: discord.TextChannel
|
|
if channel.id in blacklisted_channels:
|
|
continue
|
|
if channel.permissions_for(ctx.message.author).read_messages is False:
|
|
continue
|
|
if channel.permissions_for(ctx.guild.me).read_messages is False:
|
|
continue
|
|
channel_list.append(channel)
|
|
|
|
if len(channel_list) == 0:
|
|
return await ctx.send("There are no channels to read... This should theoretically never happen.")
|
|
|
|
embed = discord.Embed(
|
|
description="Fetching messages from the entire server this **will** take a while.",
|
|
colour=await self.bot.get_embed_colour(location=ctx.channel),
|
|
)
|
|
global_fetch_message = await ctx.send(embed=embed)
|
|
global_history = []
|
|
|
|
for channel in channel_list:
|
|
embed = discord.Embed(
|
|
title=f"Fetching messages from #{channel.name}",
|
|
description="This might take a while...",
|
|
colour=await self.bot.get_embed_colour(location=channel)
|
|
)
|
|
loading_message = await ctx.send(embed=embed)
|
|
try:
|
|
history = await self.fetch_channel_history(channel, loading_message, messages)
|
|
global_history += history
|
|
await loading_message.delete()
|
|
except discord.errors.Forbidden:
|
|
try:
|
|
await loading_message.delete()
|
|
except discord.NotFound:
|
|
continue
|
|
except discord.NotFound:
|
|
try:
|
|
await loading_message.delete()
|
|
except discord.NotFound:
|
|
continue
|
|
|
|
msg_data = self.calculate_member_perc(global_history)
|
|
# If no members are found.
|
|
if len(msg_data["users"]) == 0:
|
|
try:
|
|
await global_fetch_message.delete()
|
|
except discord.NotFound:
|
|
pass
|
|
return await ctx.send(f"Only bots have sent messages in this server... Wauw...")
|
|
|
|
top_twenty, others = self.calculate_top(msg_data)
|
|
chart = await self.create_chart(top_twenty, others, ctx.guild)
|
|
|
|
try:
|
|
await global_fetch_message.delete()
|
|
except discord.NotFound:
|
|
pass
|
|
await ctx.send(file=discord.File(chart, "chart.png"))
|
|
|
|
@checks.mod_or_permissions(manage_channels=True)
|
|
@commands.guild_only()
|
|
@commands.command()
|
|
async def ccdeny(self, ctx, channel: discord.TextChannel):
|
|
"""Add a channel to deny chatchart use."""
|
|
channel_list = await self.config.guild(ctx.guild).channel_deny()
|
|
if channel.id not in channel_list:
|
|
channel_list.append(channel.id)
|
|
await self.config.guild(ctx.guild).channel_deny.set(channel_list)
|
|
await ctx.send(f"{channel.mention} was added to the deny list for chatchart.")
|
|
|
|
@checks.mod_or_permissions(manage_channels=True)
|
|
@commands.guild_only()
|
|
@commands.command()
|
|
async def ccdenylist(self, ctx):
|
|
"""List the channels that are denied."""
|
|
no_channels_msg = "Chatchart is currently allowed everywhere in this server."
|
|
channel_list = await self.config.guild(ctx.guild).channel_deny()
|
|
if not channel_list:
|
|
msg = no_channels_msg
|
|
else:
|
|
msg = "Chatchart is not allowed in:\n"
|
|
remove_list = []
|
|
for channel in channel_list:
|
|
channel_obj = self.bot.get_channel(channel)
|
|
if not channel_obj:
|
|
remove_list.append(channel)
|
|
else:
|
|
msg += f"{channel_obj.mention}\n"
|
|
if remove_list:
|
|
new_list = [x for x in channel_list if x not in remove_list]
|
|
await self.config.guild(ctx.guild).channel_deny.set(new_list)
|
|
if len(remove_list) == len(channel_list):
|
|
msg = no_channels_msg
|
|
await ctx.send(msg)
|
|
|
|
@checks.mod_or_permissions(manage_channels=True)
|
|
@commands.guild_only()
|
|
@commands.command()
|
|
async def ccallow(self, ctx, channel: discord.TextChannel):
|
|
"""Remove a channel from the deny list to allow chatchart use."""
|
|
channel_list = await self.config.guild(ctx.guild).channel_deny()
|
|
if channel.id in channel_list:
|
|
channel_list.remove(channel.id)
|
|
else:
|
|
return await ctx.send("Channel is not on the deny list.")
|
|
await self.config.guild(ctx.guild).channel_deny.set(channel_list)
|
|
await ctx.send(f"{channel.mention} will be allowed for chatchart use.")
|
|
|
|
@checks.is_owner()
|
|
@commands.command()
|
|
async def cclimit(self, ctx, limit_amount: int = None):
|
|
"""
|
|
Limit the amount of messages someone can request.
|
|
|
|
Use `0` for no limit.
|
|
"""
|
|
if limit_amount is None:
|
|
return await ctx.send_help()
|
|
if limit_amount < 0:
|
|
return await ctx.send("You need to use a number larger than 0.")
|
|
await self.config.limit.set(limit_amount)
|
|
await ctx.send(f"Chatchart is now limited to {limit_amount} messages.")
|