From 0914cb30ccafe373c53d191bcb4782caae8b4978 Mon Sep 17 00:00:00 2001 From: Dan Date: Thu, 4 Jul 2024 17:16:27 -0400 Subject: [PATCH] FIX: Changed Profiles to use your birthday to calculate your age --- main.py | 1 + modules/user/birthday.py | 192 +++++++++++++++++++++------------------ modules/user/profiles.py | 92 ++++++++++++++++--- 3 files changed, 182 insertions(+), 103 deletions(-) diff --git a/main.py b/main.py index 11310f6..f82e873 100644 --- a/main.py +++ b/main.py @@ -99,6 +99,7 @@ class Selena(discord.Client): from modules.user.profiles import Profiles profiles = Profiles(self) profiles.setup(self.tree) + self.profiles = profiles # Properly set the profiles attribute logging.info("Profiles module loaded") diff --git a/modules/user/birthday.py b/modules/user/birthday.py index 88e204f..2777bc8 100644 --- a/modules/user/birthday.py +++ b/modules/user/birthday.py @@ -1,107 +1,119 @@ import discord -from discord import app_commands -from datetime import datetime +from discord.ext import tasks import sqlite3 -import asyncio - +import logging +import datetime class Birthday: def __init__(self, bot): self.bot = bot self.db_path = 'data/selena.db' - self._init_db() + self.logger = logging.getLogger('Birthday') + self.logger.setLevel(logging.DEBUG) + handler = logging.FileHandler(filename='log/selena.log', encoding='utf-8', mode='w') + handler.setFormatter(logging.Formatter('%(asctime)s:%(levelname)s:%(name)s:%(message)s')) + self.logger.addHandler(handler) - def _init_db(self): - with sqlite3.connect(self.db_path) as conn: - cursor = conn.cursor() - cursor.execute(''' - CREATE TABLE IF NOT EXISTS birthdays ( - user_id TEXT PRIMARY KEY, - birthday TEXT - ) - ''') - conn.commit() + self.ensure_table_exists() - def _set_birthday(self, user_id, birthday): - with sqlite3.connect(self.db_path) as conn: - cursor = conn.cursor() - cursor.execute('INSERT OR REPLACE INTO birthdays (user_id, birthday) VALUES (?, ?)', (user_id, birthday)) - conn.commit() + def ensure_table_exists(self): + conn = sqlite3.connect(self.db_path) + cursor = conn.cursor() + cursor.execute(""" + CREATE TABLE IF NOT EXISTS birthdays ( + user_id TEXT, + guild_id TEXT, + birthday TEXT, + PRIMARY KEY (user_id, guild_id) + ); + """) + conn.commit() + conn.close() + self.logger.info('Birthday table ensured in database') - def _get_birthday(self, user_id): - with sqlite3.connect(self.db_path) as conn: - cursor = conn.cursor() - cursor.execute('SELECT birthday FROM birthdays WHERE user_id = ?', (user_id,)) - row = cursor.fetchone() - return row[0] if row else None + async def set_birthday(self, user_id, guild_id, birthday): + conn = sqlite3.connect(self.db_path) + cursor = conn.cursor() + cursor.execute(""" + INSERT INTO birthdays (user_id, guild_id, birthday) + VALUES (?, ?, ?) + ON CONFLICT(user_id, guild_id) + DO UPDATE SET birthday = ? + """, (user_id, guild_id, birthday, birthday)) + conn.commit() + conn.close() + self.logger.info(f'Set birthday for user {user_id} in guild {guild_id}') - def _remove_birthday(self, user_id): - with sqlite3.connect(self.db_path) as conn: - cursor = conn.cursor() - cursor.execute('DELETE FROM birthdays WHERE user_id = ?', (user_id,)) - conn.commit() + async def get_birthday(self, user_id, guild_id): + conn = sqlite3.connect(self.db_path) + cursor = conn.cursor() + cursor.execute("SELECT birthday FROM birthdays WHERE user_id = ? AND guild_id = ?", (user_id, guild_id)) + row = cursor.fetchone() + conn.close() + return row[0] if row else None - async def set_birthday(self, interaction: discord.Interaction, date: str): - user_id = str(interaction.user.id) - try: - datetime.strptime(date, "%Y-%m-%d") - self._set_birthday(user_id, date) - await interaction.response.send_message(f'{interaction.user.mention}, your birthday has been set to {date}.', ephemeral=True) - except ValueError: - await interaction.response.send_message(f'{interaction.user.mention}, the date format is incorrect. Please use YYYY-MM-DD.', ephemeral=True) - - async def get_birthday(self, interaction: discord.Interaction): - user_id = str(interaction.user.id) - birthday = self._get_birthday(user_id) - if birthday: - await interaction.response.send_message(f'{interaction.user.mention}, your birthday is set to {birthday}.', ephemeral=True) - else: - await interaction.response.send_message(f'{interaction.user.mention}, you have not set your birthday.', ephemeral=True) - - async def remove_birthday(self, interaction: discord.Interaction): - user_id = str(interaction.user.id) - self._remove_birthday(user_id) - await interaction.response.send_message(f'{interaction.user.mention}, your birthday has been removed.', ephemeral=True) - - async def check_birthday(self): - today = datetime.today().strftime("%Y-%m-%d") - with sqlite3.connect(self.db_path) as conn: - cursor = conn.cursor() - cursor.execute('SELECT user_id FROM birthdays WHERE birthday = ?', (today,)) - users = cursor.fetchall() - - for user in users: - user_id = user[0] - user_obj = await self.bot.fetch_user(user_id) - if user_obj: - await user_obj.send(f'🎉🎂 Happy Birthday, {user_obj.mention}! 🎂🎉') - - def setup(self, tree): - @app_commands.command(name='set_birthday', description='Set your birthday') - @app_commands.describe(date='Your birthday in YYYY-MM-DD format') - async def set_birthday_command(interaction: discord.Interaction, date: str): - await self.set_birthday(interaction, date) - - @app_commands.command(name='get_birthday', description='Check your birthday') - async def get_birthday_command(interaction: discord.Interaction): - await self.get_birthday(interaction) - - @app_commands.command(name='remove_birthday', description='Remove your birthday') - async def remove_birthday_command(interaction: discord.Interaction): - await self.remove_birthday(interaction) - - tree.add_command(set_birthday_command) - tree.add_command(get_birthday_command) - tree.add_command(remove_birthday_command) + @tasks.loop(hours=24) + async def check_birthdays(self): + today = datetime.datetime.today().strftime('%Y-%m-%d') + conn = sqlite3.connect(self.db_path) + cursor = conn.cursor() + cursor.execute("SELECT user_id, guild_id FROM birthdays WHERE birthday = ?", (today,)) + rows = cursor.fetchall() + conn.close() + for user_id, guild_id in rows: + guild = self.bot.get_guild(int(guild_id)) + if guild: + user = guild.get_member(int(user_id)) + if user: + channel = guild.system_channel or next((ch for ch in guild.text_channels if ch.permissions_for(guild.me).send_messages), None) + if channel: + await channel.send(f"Happy Birthday, {user.mention}! 🎉🎂") async def setup_hook(self): - await self.bot.wait_until_ready() - while not self.bot.is_closed(): - await self.check_birthday() - await asyncio.sleep(86400) # Check once every 24 hours + self.check_birthdays.start() + self.logger.info('Started birthday check loop') + + def setup(self, tree: discord.app_commands.CommandTree): + @tree.command(name="set_birthday", description="Set your birthday") + async def set_birthday_command(interaction: discord.Interaction): + await interaction.response.send_modal(BirthdayModal(self)) + + @tree.command(name="check_birthday", description="Check your birthday") + async def check_birthday_command(interaction: discord.Interaction): + user_id = str(interaction.user.id) + guild_id = str(interaction.guild.id) + birthday = await self.get_birthday(user_id, guild_id) + if birthday: + await interaction.response.send_message(f"Your birthday is set to {birthday}.", ephemeral=True) + else: + await interaction.response.send_message("You have not set your birthday yet. Use /set_birthday to set it.", ephemeral=True) + + if not tree.get_command("set_birthday"): + tree.add_command(set_birthday_command) + + if not tree.get_command("check_birthday"): + tree.add_command(check_birthday_command) + + +class BirthdayModal(discord.ui.Modal, title="Set Birthday"): + birthday = discord.ui.TextInput(label="Your Birthday", placeholder="Enter your birthday (YYYY-MM-DD)...", required=True) + + def __init__(self, birthday_module): + super().__init__() + self.birthday_module = birthday_module + + async def on_submit(self, interaction: discord.Interaction): + user_id = str(interaction.user.id) + guild_id = str(interaction.guild.id) + try: + birthday = self.birthday.value + await self.birthday_module.set_birthday(user_id, guild_id, birthday) + await interaction.response.send_message("Your birthday has been set.", ephemeral=True) + except ValueError as e: + await interaction.response.send_message(str(e), ephemeral=True) def setup(bot): - birthday = Birthday(bot) - bot.add_cog(birthday) - bot.setup_hook = birthday.setup_hook + birthday_module = Birthday(bot) + bot.add_cog(birthday_module) + bot.birthday_module = birthday_module diff --git a/modules/user/profiles.py b/modules/user/profiles.py index d38f1c4..dae99e6 100644 --- a/modules/user/profiles.py +++ b/modules/user/profiles.py @@ -1,9 +1,11 @@ import discord import sqlite3 import logging - +from datetime import datetime class Profiles: + ALLOWED_PRONOUNS = ["he/him", "she/her", "they/them", "ask me"] + def __init__(self, bot): self.bot = bot self.db_path = 'data/selena.db' @@ -23,6 +25,7 @@ class Profiles: user_id TEXT, guild_id TEXT, pronouns TEXT, + birthday TEXT, age INTEGER, xp INTEGER DEFAULT 0, level INTEGER DEFAULT 1, @@ -40,6 +43,13 @@ class Profiles: PRIMARY KEY (user_id, guild_id, game) ); """) + + # Ensure the guild_id column exists in the game_stats table + cursor.execute("PRAGMA table_info(game_stats);") + columns = [info[1] for info in cursor.fetchall()] + if 'guild_id' not in columns: + cursor.execute("ALTER TABLE game_stats ADD COLUMN guild_id TEXT;") + conn.commit() conn.close() self.logger.info('Profile and game stats tables ensured in database') @@ -70,15 +80,20 @@ class Profiles: conn.close() self.logger.info(f'Recorded loss for user {user_id} in game {game} in guild {guild_id}') - async def update_profile(self, user_id, guild_id, pronouns=None, age=None, is_global=False): + async def update_profile(self, user_id, guild_id, pronouns=None, birthday=None, is_global=False): + if pronouns and pronouns not in self.ALLOWED_PRONOUNS: + raise ValueError("Invalid pronouns. Allowed values are: " + ", ".join(self.ALLOWED_PRONOUNS)) + age = self.calculate_age(birthday) if birthday else None + if age is not None and age < 13: + raise ValueError("You must be at least 13 years old to use Selena.") conn = sqlite3.connect(self.db_path) cursor = conn.cursor() cursor.execute(""" - INSERT INTO profiles (user_id, guild_id, pronouns, age, is_global) - VALUES (?, ?, ?, ?, ?) + INSERT INTO profiles (user_id, guild_id, pronouns, birthday, age, is_global) + VALUES (?, ?, ?, ?, ?, ?) ON CONFLICT(user_id, guild_id) - DO UPDATE SET pronouns = COALESCE(?, pronouns), age = COALESCE(?, age), is_global = ? - """, (user_id, guild_id, pronouns, age, is_global, pronouns, age, is_global)) + DO UPDATE SET pronouns = COALESCE(?, pronouns), birthday = COALESCE(?, birthday), age = COALESCE(?, age), is_global = ? + """, (user_id, guild_id, pronouns, birthday, age, is_global, pronouns, birthday, age, is_global)) conn.commit() conn.close() self.logger.info(f'Updated profile for user {user_id} in guild {guild_id}') @@ -87,7 +102,7 @@ class Profiles: conn = sqlite3.connect(self.db_path) cursor = conn.cursor() cursor.execute(""" - SELECT pronouns, age, xp, level, is_global FROM profiles + SELECT pronouns, birthday, age, xp, level, is_global FROM profiles WHERE user_id = ? AND (guild_id = ? OR is_global = 1) """, (user_id, guild_id)) row = cursor.fetchone() @@ -102,13 +117,18 @@ class Profiles: conn.close() return rows + def calculate_age(self, birthday): + if not birthday: + return None + today = datetime.today() + birthdate = datetime.strptime(birthday, "%Y-%m-%d") + age = today.year - birthdate.year - ((today.month, today.day) < (birthdate.month, birthdate.day)) + return age + def setup(self, tree: discord.app_commands.CommandTree): @tree.command(name="set_profile", description="Set your profile information") - async def set_profile_command(interaction: discord.Interaction, pronouns: str = None, age: int = None, is_global: bool = False): - user_id = str(interaction.user.id) - guild_id = str(interaction.guild.id) - await self.update_profile(user_id, guild_id, pronouns, age, is_global) - await interaction.response.send_message("Your profile has been updated.", ephemeral=True) + async def set_profile_command(interaction: discord.Interaction): + await interaction.response.send_message("Please select your pronouns and set your birthday:", view=ProfileView(self)) @tree.command(name="profile", description="View your profile") async def profile_command(interaction: discord.Interaction): @@ -116,11 +136,12 @@ class Profiles: guild_id = str(interaction.guild.id) profile = await self.get_profile(user_id, guild_id) if profile: - pronouns, age, xp, level, is_global = profile + pronouns, birthday, age, xp, level, is_global = profile game_stats = await self.get_game_stats(user_id, guild_id) games_info = "\n".join([f"{game}: {wins}W/{losses}L" for game, wins, losses in game_stats]) embed = discord.Embed(title=f"{interaction.user.display_name}'s Profile") embed.add_field(name="Pronouns", value=pronouns or "Not set", inline=True) + embed.add_field(name="Birthday", value=birthday or "Not set", inline=True) embed.add_field(name="Age", value=age or "Not set", inline=True) embed.add_field(name="XP", value=xp, inline=True) embed.add_field(name="Level", value=level, inline=True) @@ -137,6 +158,51 @@ class Profiles: tree.add_command(profile_command) +class ProfileView(discord.ui.View): + def __init__(self, profiles): + super().__init__() + self.profiles = profiles + self.add_item(PronounSelect(profiles)) + + @discord.ui.button(label="Set Birthday", style=discord.ButtonStyle.primary) + async def set_birthday_button(self, interaction: discord.Interaction, button: discord.ui.Button): + await interaction.response.send_modal(BirthdayModal(self.profiles)) + + +class PronounSelect(discord.ui.Select): + def __init__(self, profiles): + self.profiles = profiles + options = [discord.SelectOption(label=pronoun, value=pronoun) for pronoun in Profiles.ALLOWED_PRONOUNS] + super().__init__(placeholder="Select your pronouns...", min_values=1, max_values=1, options=options) + + async def callback(self, interaction: discord.Interaction): + user_id = str(interaction.user.id) + guild_id = str(interaction.guild.id) + try: + await self.profiles.update_profile(user_id, guild_id, pronouns=self.values[0]) + await interaction.response.send_message(f"Your pronouns have been set to {self.values[0]}.", ephemeral=True) + except ValueError as e: + await interaction.response.send_message(str(e), ephemeral=True) + + +class BirthdayModal(discord.ui.Modal, title="Set Birthday"): + birthday = discord.ui.TextInput(label="Your Birthday", placeholder="Enter your birthday (YYYY-MM-DD)...", required=True) + + def __init__(self, profiles): + super().__init__() + self.profiles = profiles + + async def on_submit(self, interaction: discord.Interaction): + user_id = str(interaction.user.id) + guild_id = str(interaction.guild.id) + try: + birthday = self.birthday.value + await self.profiles.update_profile(user_id, guild_id, birthday=birthday) + await interaction.response.send_message("Your birthday has been set.", ephemeral=True) + except ValueError as e: + await interaction.response.send_message(str(e), ephemeral=True) + + def setup(bot): profiles = Profiles(bot) profiles.setup(bot.tree)