From f0a25b78bce0814881b9f817c5f98e3438400deb Mon Sep 17 00:00:00 2001 From: LBlend Date: Tue, 19 Dec 2023 19:41:52 +0100 Subject: [PATCH] add mbti feature --- Dockerfile | 2 +- requirements.txt | 1 + src/cogs/user_facts.py | 542 ++++++++++++++++++++++++++++------------- 3 files changed, 375 insertions(+), 170 deletions(-) diff --git a/Dockerfile b/Dockerfile index 58c2799..ba5a266 100644 --- a/Dockerfile +++ b/Dockerfile @@ -6,7 +6,7 @@ COPY requirements.txt . RUN python3 -m pip install -r requirements.txt -RUN apt update && apt install -y imagemagick ffmpeg libsm6 libxext6 +RUN apt update && apt install -y imagemagick ffmpeg libsm6 libxext6 graphviz COPY . . diff --git a/requirements.txt b/requirements.txt index b01a166..68db345 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,6 +1,7 @@ aiocron==1.8.* asyncio==3.4.* discord.py==2.3.* +graphviz==0.20.* mcrcon==0.7.* moviepy==1.0.* nltk==3.8.* diff --git a/src/cogs/user_facts.py b/src/cogs/user_facts.py index 86681ff..c2d335a 100644 --- a/src/cogs/user_facts.py +++ b/src/cogs/user_facts.py @@ -1,170 +1,374 @@ -import discord -from discord import app_commands -from discord.ext import commands - -from cogs.utils import discord_utils -from cogs.utils import embed_templates -from cogs.utils import misc_utils - - -class UserFacts(commands.Cog): - def __init__(self, bot): - self.bot = bot - self.cursor = self.bot.db_connection.cursor() - self.init_db() - - def init_db(self): - """Create the necessary tables for the birthday cog to work""" - - self.cursor.execute( - """ - CREATE TABLE IF NOT EXISTS user_facts ( - discord_id BIGINT PRIMARY KEY, - mbti CHAR(4), - height INT - ); - """ - ) - - height_group = app_commands.Group(name="høyde", description="Se, endre eller fjern høyde for brukere på serveren") - - @height_group.command(name="se", description="Se høyden til en bruker") - async def height_see(self, interaction: discord.Interaction, bruker: discord.Member = None): - """ - See the height of a user - - Parameters - ---------- - interaction (discord.Interaction): The interaction object - bruker (discord.Member): The user to see the height of - """ - - if not bruker: - bruker = interaction.user - - self.cursor.execute( - """ - SELECT height FROM user_facts - WHERE discord_id = %s and height IS NOT NULL; - """, - (bruker.id,), - ) - height = self.cursor.fetchone() - - if not height: - return await interaction.response.send_message( - embed=embed_templates.error_warning(interaction, text="Brukeren har ikke lagt inn høyden sin") - ) - - height = height[0] - - inches = height * (1 / 2.54) - feet = int(inches * (1 / 12)) - inches = int(inches - (feet * 12)) - - embed = discord.Embed(color=bruker.color) - embed.add_field(name="Høyde", value=f"{height} cm") - embed.add_field(name="Høyde (🦅🦅🇺🇸🦅🦅)", value=f"{feet}'{inches}\"") - embed.set_author(name=bruker.global_name, icon_url=bruker.avatar) - - await interaction.response.send_message(embed=embed) - - @height_group.command(name="sett", description="Sett høyden din") - @app_commands.rename(height_cm="høyde_cm") - async def height_set(self, interaction: discord.Interaction, height_cm: app_commands.Range[int, 50, 250]): - """ - Set your height - - Parameters - ---------- - interaction (discord.Interaction): The interaction object - height_cm (int): The height of the user - """ - - self.cursor.execute( - """ - INSERT INTO user_facts (discord_id, height) - VALUES (%s, %s) - ON CONFLICT (discord_id) DO UPDATE - SET height = %s; - """, - (interaction.user.id, height_cm, height_cm), - ) - - await interaction.response.send_message(embed=embed_templates.success(interaction, text="Høyde satt!")) - - @height_group.command(name="fjern", description="Fjern høyden din") - async def height_remove(self, interaction: discord.Interaction): - """ - Remove your height - - Parameters - ---------- - interaction (discord.Interaction): The interaction object - """ - - self.cursor.execute( - """ - UPDATE user_facts - SET height = NULL - WHERE discord_id = %s; - """, - (interaction.user.id,), - ) - - if self.cursor.rowcount == 0: - return await interaction.response.send_message( - embed=embed_templates.success( - interaction, text="Du hadde ikke lagt inn høyden din fra før av men ok :)" - ) - ) - - await interaction.response.send_message(embed=embed_templates.success(interaction, text="Høyde fjernet!")) - - @height_group.command(name="leaderboard", description="Se en leaderboard over høyder") - async def height_leaderboard(self, interaction: discord.Interaction): - """ - See a leaderboard of heights - - Parameters - ---------- - interaction (discord.Interaction): The interaction object - """ - - self.cursor.execute( - """ - SELECT discord_id, height FROM user_facts +import discord +import graphviz +import numpy as np +from discord import app_commands +from discord.ext import commands + +from cogs.utils import discord_utils +from cogs.utils import embed_templates +from cogs.utils import misc_utils + + +class UserFacts(commands.Cog): + def __init__(self, bot): + self.bot = bot + self.cursor = self.bot.db_connection.cursor() + self.init_db() + + self.mbti_codes = { + "INTJ", + "INTP", + "ENTJ", + "ENTP", + "INFJ", + "INFP", + "ENFJ", + "ENFP", + "ISTJ", + "ISFJ", + "ESTJ", + "ESFJ", + "ISTP", + "ISFP", + "ESTP", + "ESFP", + } + self.mbti_list = list(self.mbti_codes) + self.similarity_matrix = self._create_similarity_matrix() + + def init_db(self): + """Create the necessary tables for the birthday cog to work""" + + self.cursor.execute( + """ + CREATE TABLE IF NOT EXISTS user_facts ( + discord_id BIGINT PRIMARY KEY, + mbti CHAR(4), + height INT + ); + """ + ) + + height_group = app_commands.Group(name="høyde", description="Se, endre eller fjern høyde for brukere på serveren") + + @height_group.command(name="se", description="Se høyden til en bruker") + async def height_see(self, interaction: discord.Interaction, bruker: discord.Member = None): + """ + See the height of a user + + Parameters + ---------- + interaction (discord.Interaction): The interaction object + bruker (discord.Member): The user to see the height of + """ + + if not bruker: + bruker = interaction.user + + self.cursor.execute( + """ + SELECT height FROM user_facts + WHERE discord_id = %s and height IS NOT NULL; + """, + (bruker.id,), + ) + height = self.cursor.fetchone() + + if not height: + return await interaction.response.send_message( + embed=embed_templates.error_warning(interaction, text="Brukeren har ikke lagt inn høyden sin") + ) + + height = height[0] + + inches = height * (1 / 2.54) + feet = int(inches * (1 / 12)) + inches = int(inches - (feet * 12)) + + embed = discord.Embed(color=bruker.color) + embed.add_field(name="Høyde", value=f"{height} cm") + embed.add_field(name="Høyde (🦅🦅🇺🇸🦅🦅)", value=f"{feet}'{inches}\"") + embed.set_author(name=bruker.global_name, icon_url=bruker.avatar) + + await interaction.response.send_message(embed=embed) + + @height_group.command(name="sett", description="Sett høyden din") + @app_commands.rename(height_cm="høyde_cm") + async def height_set(self, interaction: discord.Interaction, height_cm: app_commands.Range[int, 50, 250]): + """ + Set your height + + Parameters + ---------- + interaction (discord.Interaction): The interaction object + height_cm (int): The height of the user + """ + + self.cursor.execute( + """ + INSERT INTO user_facts (discord_id, height) + VALUES (%s, %s) + ON CONFLICT (discord_id) DO UPDATE + SET height = %s; + """, + (interaction.user.id, height_cm, height_cm), + ) + + await interaction.response.send_message(embed=embed_templates.success(interaction, text="Høyde satt!")) + + @height_group.command(name="fjern", description="Fjern høyden din") + async def height_remove(self, interaction: discord.Interaction): + """ + Remove your height + + Parameters + ---------- + interaction (discord.Interaction): The interaction object + """ + + self.cursor.execute( + """ + UPDATE user_facts + SET height = NULL + WHERE discord_id = %s; + """, + (interaction.user.id,), + ) + + if self.cursor.rowcount == 0: + return await interaction.response.send_message( + embed=embed_templates.success( + interaction, text="Du hadde ikke lagt inn høyden din fra før av men ok :)" + ) + ) + + await interaction.response.send_message(embed=embed_templates.success(interaction, text="Høyde fjernet!")) + + @height_group.command(name="leaderboard", description="Se en leaderboard over høyder") + async def height_leaderboard(self, interaction: discord.Interaction): + """ + See a leaderboard of heights + + Parameters + ---------- + interaction (discord.Interaction): The interaction object + """ + + self.cursor.execute( + """ + SELECT discord_id, height FROM user_facts WHERE height IS NOT NULL - ORDER BY height DESC; - """ - ) - - if not (result := self.cursor.fetchall()): - return await interaction.response.send_message( - embed=embed_templates.error_warning(interaction, text="Ingen har lagt inn høyden sin enda") - ) - - result_formatted = list( - map( - lambda s: f"**#{s[0]+1}** <@{s[1][0]}> - `{s[1][1]}` cm", - enumerate(result), - ) - ) - - paginator = misc_utils.Paginator(result_formatted) - view = discord_utils.Scroller(paginator, interaction.user) - - embed = view.construct_embed(discord.Embed(title="Våre høyeste og laveste (👑)")) - await interaction.response.send_message(embed=embed, view=view) - - -async def setup(bot: commands.Bot): - """ - Add the cog to the bot on extension load - - Parameters - ---------- - bot (commands.Bot): Bot instance - """ - - await bot.add_cog(UserFacts(bot)) + ORDER BY height DESC; + """ + ) + + if not (result := self.cursor.fetchall()): + return await interaction.response.send_message( + embed=embed_templates.error_warning(interaction, text="Ingen har lagt inn høyden sin enda") + ) + + result_formatted = list( + map( + lambda s: f"**#{s[0]+1}** <@{s[1][0]}> - `{s[1][1]}` cm", + enumerate(result), + ) + ) + + paginator = misc_utils.Paginator(result_formatted) + view = discord_utils.Scroller(paginator, interaction.user) + + embed = view.construct_embed(discord.Embed(title="Våre høyeste og laveste (👑)")) + await interaction.response.send_message(embed=embed, view=view) + + mbti_group = app_commands.Group(name="mbti", description="Se, endre eller fjern MBTI for brukere på serveren") + + @mbti_group.command(name="se", description="Se MBTI-en til en bruker") + async def mbti_see(self, interaction: discord.Interaction, bruker: discord.Member = None): + """ + See the MBTI of a user + + Parameters + ---------- + interaction (discord.Interaction): The interaction object + bruker (discord.Member): The user to see the MBTI of + """ + + if not bruker: + bruker = interaction.user + + self.cursor.execute( + """ + SELECT mbti FROM user_facts + WHERE discord_id = %s AND mbti IS NOT NULL; + """, + (bruker.id,), + ) + user_mbti = self.cursor.fetchone() + + if not user_mbti: + return await interaction.response.send_message( + embed=embed_templates.error_warning(interaction, text="Brukeren har ikke lagt inn MBTI-en sin") + ) + + user_mbti = user_mbti[0] + + embed = discord.Embed(color=bruker.color, title="MBTI", description=user_mbti) + embed.set_author(name=bruker.global_name, icon_url=bruker.avatar) + + self.cursor.execute( + """ + SELECT discord_id, mbti FROM user_facts + WHERE mbti IS NOT NULL and discord_id != %s; + """, + (bruker.id,), + ) + + if not (results := self.cursor.fetchall()): + return await interaction.response.send_message(embed=embed) + + others = [] + for discord_id, mbti in results: + if user := self.bot.get_user(discord_id): + others.append((user, mbti)) + elif user := await interaction.guild.fetch_member(discord_id): + others.append((user, mbti)) + else: + others.append((discord_id, mbti)) + + self._create_mbti_graph((bruker, user_mbti), others) + + with open(f"src/assets/temp/{bruker.id}_mbti.png", "rb") as f: + image = discord.File(f, filename=f"{bruker.id}_mbti.png") + + embed.set_image(url=f"attachment://{bruker.id}_mbti.png") + await interaction.response.send_message(embed=embed, file=image) + + @mbti_group.command(name="sett", description="Sett MBTI-en din") + async def mbti_set(self, interaction: discord.Interaction, mbti: str): + """ + Set your MBTI + + Parameters + ---------- + interaction (discord.Interaction): The interaction object + mbti (str): The MBTI of the user + """ + + if mbti.upper() not in self.mbti_codes: + return await interaction.response.send_message( + embed=embed_templates.error_warning(interaction, text="Ugyldig MBTI") + ) + + self.cursor.execute( + """ + INSERT INTO user_facts (discord_id, mbti) + VALUES (%s, %s) + ON CONFLICT (discord_id) DO UPDATE + SET mbti = %s; + """, + (interaction.user.id, mbti.upper(), mbti.upper()), + ) + + await interaction.response.send_message(embed=embed_templates.success(interaction, text="MBTI satt!")) + + @mbti_set.autocomplete("mbti") + async def mbti_set_autocomplete_callback(self, interaction: discord.Interaction, current: str): + """ + Autocomplete for the MBTI set command + + Parameters + ---------- + interaction (discord.Interaction): The interaction object + current (str): The current input + """ + + return [ + app_commands.Choice(name=mbti, value=mbti) for mbti in self.mbti_codes if mbti.startswith(current.upper()) + ] + + @mbti_group.command(name="fjern", description="Fjern MBTI-en din") + async def mbti_remove(self, interaction: discord.Interaction): + """ + Remove your MBTI + + Parameters + ---------- + interaction (discord.Interaction): The interaction object + """ + + self.cursor.execute( + """ + UPDATE user_facts + SET mbti = NULL + WHERE discord_id = %s; + """, + (interaction.user.id,), + ) + + if self.cursor.rowcount == 0: + return await interaction.response.send_message( + embed=embed_templates.success( + interaction, text="Du hadde ikke lagt inn MBTIen din fra før av men ok :)" + ) + ) + + await interaction.response.send_message(embed=embed_templates.success(interaction, text="MBTI fjernet!")) + + def _create_similarity_matrix(self): + def similarity(mbti_1, mbti_2): + similarity = 0 + for i in range(4): + if mbti_1[i] == mbti_2[i]: + similarity += 1 + return similarity + + # Create an empty similarity matrix + similarity_matrix = np.zeros((len(self.mbti_list), len(self.mbti_list))) + + # Fill the similarity matrix + for i, mbti in enumerate(self.mbti_list): + for j, mbti2 in enumerate(self.mbti_list): + similarity_matrix[i, j] = similarity(mbti, mbti2) + + return similarity_matrix + + def _create_mbti_graph(self, user_mbti: tuple[discord.Member, str], others: list[tuple[discord.Member, str]]): + user, mbti = user_mbti + + # Create a Graphviz graph + graph = graphviz.Graph(engine="neato", format="png") + graph.attr("graph", overlap="false") + graph.node_attr["style"] = "filled" + graph.node_attr["shape"] = "circle" + + # Invoking user's node + graph.node(str(user.global_name), f"{user.global_name}\n{mbti}", color="green") + + user_index = self.mbti_list.index(mbti) + + # Add nodes and edges to other users + for other_index, other in enumerate(others): + weight = (1 / self.similarity_matrix[user_index, other_index]) * 6 + if weight == float("inf"): + weight = 6 + elif weight == 0: + weight = 6.5 + + other_user, other_mbti = other + + graph.node(other_user.global_name, f"{other_user.global_name}\n{other_mbti}") + graph.edge(user.global_name, other_user.global_name, len=str(weight), weight=str(weight)) + + graph.render(f"{user.id}_mbti", directory="./src/assets/temp", view=False, overwrite_source=True) + + +async def setup(bot: commands.Bot): + """ + Add the cog to the bot on extension load + + Parameters + ---------- + bot (commands.Bot): Bot instance + """ + + await bot.add_cog(UserFacts(bot))