From 23c4047bf731466fd03d5adf92104769b29b2075 Mon Sep 17 00:00:00 2001 From: Nebulo9 Date: Sun, 19 Nov 2023 13:47:16 +0100 Subject: [PATCH] Added productivity feature --- .gitignore | 3 +- main.py | 71 ++++++++++++++-------- modules/cogs/birthday.py | 113 +++++++++++++++++++---------------- modules/cogs/productivity.py | 107 +++++++++++++++++++++++++++++++++ modules/setup/bot.py | 15 ++++- modules/setup/data.py | 26 +++++++- modules/setup/logger.py | 2 +- 7 files changed, 254 insertions(+), 83 deletions(-) create mode 100644 modules/cogs/productivity.py diff --git a/.gitignore b/.gitignore index e31e7d5..cfbe543 100644 --- a/.gitignore +++ b/.gitignore @@ -159,4 +159,5 @@ cython_debug/ # option (not recommended) you can uncomment the following to ignore the entire idea folder. #.idea/ -data/ \ No newline at end of file +data/ +.vscode/ \ No newline at end of file diff --git a/main.py b/main.py index b8b93d9..d2241b2 100644 --- a/main.py +++ b/main.py @@ -1,11 +1,14 @@ import os, discord, json, sys from argparse import ArgumentParser from datetime import datetime +from asyncio import create_task from dotenv import load_dotenv +from discord import option from discord.ext import tasks, commands -from modules.setup.bot import bot +from modules.setup.bot import bot, reload_feature from modules.setup.logger import LOGGER -from modules.setup.data import DATA_DIR, get_guild_data +from modules.setup.data import get_guild_data, create_guild_data, delete_guild_data, save_guild_data +from modules.cogs.productivity import TASKS as PRODUCTIVITY_TASKS, send_reminder PARSER = ArgumentParser(description='BotMafieux for Discord.') PARSER.add_argument('--guild_id',type=int,help='The guild ID.') @@ -45,33 +48,35 @@ async def help(ctx:discord.ApplicationContext): await ctx.send_response(embed=embed,ephemeral=True) @bot.slash_command(name='reload',description='Reloads the bot.') +@option(name='feature',description='The feature to reload.',required=True,choices=['birthday','productivity']) @commands.is_owner() -async def reload(ctx:discord.ApplicationContext): +async def reload(ctx:discord.ApplicationContext, feature:str): """Reloads the bot.""" command_name = 'reload' guild = ctx.guild channel = ctx.channel - LOGGER.debug(f'{ctx.author.name} used /{command_name}.') - await ctx.send_response('Reloading...',ephemeral=True) - args = [f'--guild_id={guild.id}',f'--channel_id={channel.id}','--reload'] - os.execl(sys.executable, sys.executable, __file__, *args) + LOGGER.debug(f'{ctx.author.name} used /{command_name} {feature}.') + try: + reload_feature(feature) + await ctx.send_response(f'{feature} reloaded!',ephemeral=True) + except Exception as e: + LOGGER.error(f'Error reloading feature {feature}: {e}') + await ctx.send_response(f'Error reloading feature {feature}: {e}',ephemeral=True) + @bot.event async def on_guild_join(guild:discord.Guild): - LOGGER.info(f'Bot joined guild {guild.name} ({guild.id}).') - path = os.path.join(DATA_DIR, f'{guild.id}.json') - if not os.path.exists(path): - LOGGER.info(f'Creating data file for guild {guild.id}') - with open(path, 'x') as f: - json.dump({}, f, indent=2) + guild_id = guild.id + guild_name = guild.name + LOGGER.info(f'Bot joined guild {guild_name} ({guild_id}).') + create_guild_data(guild_id) @bot.event async def on_guild_remove(guild:discord.Guild): - LOGGER.info(f'Bot left guild {guild.name} ({guild.id}).') - path = os.path.join(DATA_DIR, f'{guild.id}.json') - if os.path.exists(path): - LOGGER.info(f'Deleting data file for guild {guild.id}.') - os.remove(path) + guild_id = guild.id + guild_name = guild.name + LOGGER.info(f'Bot left guild {guild_name} ({guild_id}).') + delete_guild_data(guild_id) @bot.event async def on_ready(): @@ -81,15 +86,29 @@ async def on_ready(): guild = bot.get_guild(EXEC_ARGS.guild_id) channel = guild.get_channel_or_thread(EXEC_ARGS.channel_id) await channel.send(f'{bot.user.name} Bot reloaded!',silent=True) - # Creates guild data files if they don't exist - for guild in bot.guilds: - guild_id = guild.id - path = os.path.join(DATA_DIR, f'{guild_id}.json') - if not os.path.exists(path): - LOGGER.info(f'Creating data file for guild {guild_id}') - with open(path, 'x') as f: - json.dump({'features': {}}, f, indent=2) birthday_anouncements_task.start() + for guild in bot.guilds: + guild_data = get_guild_data(guild.id) + if 'features' in guild_data.keys(): + if 'productivity' in guild_data['features'].keys(): + if 'reminders' in guild_data['features']['productivity'].keys(): + for user_id in guild_data['features']['productivity']['reminders'].keys(): + user_data = guild_data['features']['productivity']['reminders'][user_id] + if user_data['enabled']: + next_reminder = user_data['next_reminder'] + if next_reminder > datetime.now().timestamp(): + cooldown = next_reminder - datetime.now().timestamp() + user = guild.get_member(int(user_id)) + channel = guild.get_channel(user_data['channel']) + custom_message = user_data['custom_message'] + PRODUCTIVITY_TASKS[user_id] = create_task(send_reminder(cooldown,user,channel,custom_message)) + else: + user_data['next_reminder'] = datetime.now().timestamp() + user_data['cooldown'] + save_guild_data(guild.id, guild_data) + user = guild.get_member(int(user_id)) + channel = guild.get_channel(user_data['channel']) + custom_message = user_data['custom_message'] + PRODUCTIVITY_TASKS[user_id] = create_task(send_reminder(user_data['cooldown'],user,channel,custom_message)) if __name__ == '__main__': load_dotenv() diff --git a/modules/cogs/birthday.py b/modules/cogs/birthday.py index fe62835..f3c718d 100644 --- a/modules/cogs/birthday.py +++ b/modules/cogs/birthday.py @@ -2,7 +2,7 @@ import discord, re from discord import option from discord.ext import commands from ..setup.logger import LOGGER -from ..setup.data import get_guild_data, save_guild_data +from ..setup.data import get_guild_data, save_guild_data, is_feature_enabled class BirthdayCog(commands.Cog): def __init__(self, bot): @@ -17,43 +17,40 @@ class BirthdayCog(commands.Cog): guild_id = ctx.guild.id guild_data = get_guild_data(guild_id) author = ctx.author - if re.match(r'\d{2}\/\d{2}',date): # Check if date is in DAY/MONTH format - if for_user: - if author.guild_permissions.administrator: # Check if author is an administrator in case they want to set the birthday for another user - LOGGER.debug(f'{author.name} used /{command_name} {date} for {for_user.name}.') - if 'birthday' not in guild_data['features'].keys(): - guild_data['features']['birthday'] = dict() - if 'birthdays' not in guild_data['features']['birthday'].keys(): - guild_data['features']['birthday']['birthdays'] = dict() - if str(for_user.id) not in guild_data['features']['birthday']['birthdays'].keys(): - guild_data['features']['birthday'][str(for_user.id)] = dict() - guild_data['features']['birthday']['birthdays'][str(for_user.id)]['date'] = date - guild_data['features']['birthday']['birthdays'][str(for_user.id)]['announcements'] = True + if is_feature_enabled('birthday',data=guild_data): + if re.match(r'\d{2}\/\d{2}',date): # Check if date is in DAY/MONTH format + feature_data = guild_data['features']['birthday'] + if 'birthdays' not in feature_data.keys(): + feature_data['birthdays'] = dict() + if str(for_user.id) not in feature_data['birthdays'].keys(): + feature_data[str(for_user.id)] = dict() + if for_user: + if author.guild_permissions.administrator: # Check if author is an administrator in case they want to set the birthday for another user + LOGGER.debug(f'{author.name} used /{command_name} {date} for {for_user.name}.') + feature_data['birthdays'][str(for_user.id)]['date'] = date + feature_data['birthdays'][str(for_user.id)]['announcements'] = True - save_guild_data(guild_id, guild_data) - await ctx.send_response(f'Birthday for {for_user.name} has been set to {date}.',ephemeral=True) + save_guild_data(guild_id, guild_data) + await ctx.send_response(f'Birthday for {for_user.name} has been set to {date}.',ephemeral=True) + else: + LOGGER.debug(f'{author.name} tried to use /{command_name} {date} for {for_user.name} but is not an administrator.') + await ctx.send_response('You must be an administrator to run the command with "for_user".',ephemeral=True) else: - LOGGER.debug(f'{author.name} tried to use /{command_name} {date} for {for_user.name} but is not an administrator.') - await ctx.send_response('You must be an administrator to run the command with "for_user".',ephemeral=True) + # Set birthday for author + LOGGER.debug(f'{author.name} used /{command_name} {date}.') + feature_data['birthdays'][str(author.id)]['date'] = date + feature_data['birthdays'][str(author.id)]['announcements'] = True + save_guild_data(guild_id, guild_data) + + # The status defines the human readable status of the announcements to be displayed in the response. + announcements_status = 'will be' if feature_data['birthdays'][str(author.id)]['announcements'] else 'will not be' + await ctx.send_response(f'Your birthday has been set to {date} and {announcements_status} announced.',ephemeral=True) else: - # Set birthday for author - LOGGER.debug(f'{author.name} used /{command_name} {date}.') - if 'birthday' not in guild_data['features'].keys(): - guild_data['features']['birthday'] = dict() - if 'birthdays' not in guild_data['features']['birthday'].keys(): - guild_data['features']['birthday']['birthdays'] = dict() - if author.id not in guild_data['features']['birthday']['birthdays'].keys(): - guild_data['features']['birthday']['birthdays'][str(author.id)] = dict() - guild_data['features']['birthday']['birthdays'][str(author.id)]['date'] = date - guild_data['features']['birthday']['birthdays'][str(author.id)]['announcements'] = True - save_guild_data(guild_id, guild_data) - - # The status defines the human readable status of the announcements to be displayed in the response. - announcements_status = 'will be' if guild_data['features']['birthday']['birthdays'][str(author.id)]['announcements'] else 'will not be' - await ctx.send_response(f'{author.mention} has set their birthday to {date} and {announcements_status} announced.',ephemeral=True) + LOGGER.debug(f'{author.name} used /set_birthday {date} but the format is not correct.') + await ctx.send_response('The date must is DAY/MONTH format.',ephemeral=True) else: - LOGGER.debug(f'{author.name} used /set_birthday {date} but the format is not correct.') - await ctx.send_response('The date must is DAY/MONTH format.',ephemeral=True) + LOGGER.debug(f'{author.name} tried to use /set_birthday {date} but the birthday feature is not enabled.') + await ctx.send_response('The birthday feature is not enabled.',ephemeral=True) @commands.slash_command(description='Enable or disable birthday announcements.') @option(name='enable',description='Enable birthday announcements.',required=True) @@ -64,20 +61,25 @@ class BirthdayCog(commands.Cog): guild_id = ctx.guild.id guild_data = get_guild_data(guild_id) or dict() author = ctx.author - if for_user: - if author.guild_permissions.administrator: # Check if author is an administrator in case they want to set the birthday for another user - LOGGER.debug(f'{author.name} used /{command_name} {enable} for {for_user.name}.') - guild_data['features']['birthday']['birthdays'][str(for_user.id)]['announcements'] = enable - save_guild_data(guild_id, guild_data) - await ctx.send_response(f'Announcements for {for_user.name} are set to {enable}',ephemeral=True) + if is_feature_enabled('birthday',data=guild_data): + feature_data = guild_data['features']['birthday'] + if for_user: + if author.guild_permissions.administrator: # Check if author is an administrator in case they want to set the birthday for another user + LOGGER.debug(f'{author.name} used /{command_name} {enable} for {for_user.name}.') + feature_data[str(for_user.id)]['announcements'] = enable + save_guild_data(guild_id, guild_data) + await ctx.send_response(f'Announcements for {for_user.name} are set to {enable}',ephemeral=True) + else: + LOGGER.debug(f'{author.name} tried to use /{command_name} {enable} for {for_user.name} but is not an administrator.') + await ctx.send_response('You must be an administrator to run the command with "for_user".',ephemeral=True) else: - LOGGER.debug(f'{author.name} tried to use /{command_name} {enable} for {for_user.name} but is not an administrator.') - await ctx.send_response('You must be an administrator to run the command with "for_user".',ephemeral=True) + LOGGER.debug(f'{author.name} used /{command_name} {enable}.') + feature_data['birthdays'][str(author.id)]['announcements'] = enable + save_guild_data(guild_id, guild_data) + await ctx.send_response(f'Your birthday announcements have been set to {enable}',ephemeral=True) else: - LOGGER.debug(f'{author.name} used /{command_name} {enable}.') - guild_data['features']['birthday']['birthdays'][str(author.id)]['announcements'] = enable - save_guild_data(guild_id, guild_data) - await ctx.send_response(f'Your birthday announcements have been set to {enable}',ephemeral=True) + LOGGER.debug(f'{author.name} tried to use /{command_name} {enable} but the birthday feature is not enabled.') + await ctx.send_response('The birthday feature is not enabled.',ephemeral=True) @commands.slash_command(description='Sets the channel to send birthday announcements to.') @option(name='channel',description='The channel to send birthday announcements to.',required=True,type=discord.TextChannel) @@ -87,12 +89,17 @@ class BirthdayCog(commands.Cog): guild_id = ctx.guild.id guild_data = get_guild_data(guild_id) or dict() author = ctx.author - if author.guild_permissions.administrator: - LOGGER.debug(f'{author.name} used /{command_name} {channel.name}.') - guild_data['features']['birthday']['birthday_announcements_channel'] = channel.id - save_guild_data(guild_id, guild_data) - await ctx.send_response(f'Birthday announcements will be sent to {channel.mention}.',ephemeral=True) + if is_feature_enabled('birthday',data=guild_data): + feature_data = guild_data['features']['birthday'] + if author.guild_permissions.administrator: + LOGGER.debug(f'{author.name} used /{command_name} {channel.name}.') + feature_data['birthday_announcements_channel'] = channel.id + save_guild_data(guild_id, guild_data) + await ctx.send_response(f'Birthday announcements will be sent to {channel.mention}.',ephemeral=True) + else: + LOGGER.debug(f'{author.name} tried to use /{command_name} {channel.mention} but is not an administrator.') + await ctx.send_response('You must be an administrator to run this command.',ephemeral=True) else: - LOGGER.debug(f'{author.name} tried to use /{command_name} {channel.mention} but is not an administrator.') - await ctx.send_response('You must be an administrator to run this command.',ephemeral=True) + LOGGER.debug(f'{author.name} tried to use /{command_name} {channel.mention} but the birthday feature is not enabled.') + await ctx.send_response('The birthday feature is not enabled.',ephemeral=True) diff --git a/modules/cogs/productivity.py b/modules/cogs/productivity.py new file mode 100644 index 0000000..d705060 --- /dev/null +++ b/modules/cogs/productivity.py @@ -0,0 +1,107 @@ +import discord +from datetime import datetime +from discord.ext import commands, tasks +from discord import option +from typing import Union +from asyncio import sleep, create_task +from ..setup.logger import LOGGER +from ..setup.data import get_guild_data, save_guild_data, is_feature_enabled + +TASKS = {} + +async def send_reminder(cooldown:int,author:discord.Member,channel:discord.TextChannel,reminder_message:str=None): + while True: + await sleep(cooldown) + LOGGER.debug(f'Sending productivity reminder to {author.name} in {channel.name} of {channel.guild.name}.') + if reminder_message: + await channel.send(f'{author.mention} {reminder_message}') + else: + await channel.send(f'{author.mention} You have not been active in this channel for a while. Please consider being more productive.') + +class ProductivityCog(commands.Cog): + + def __init__(self,bot:discord.Bot) -> None: + self.bot = bot + + @commands.slash_command(description='Configurate productivity reminder.') + @option(name='channel',description='The channel to send productivity reminders in.',required=True) + @option(name='enable',description='Enable or disable productivity reminders.',required=True) + @option(name='days',description='Number of days to wait before sending a reminder if the user has not been active in the channel.',required=True,type=int) + @option(name='custom_message',description='Custom message to send with the reminder.',required=False,type=str) + @option(name='for_user',description='The user to configurate productivity reminders for.',required=False,type=discord.Member) + async def productivity_config(self,ctx:discord.ApplicationContext, channel:Union[discord.TextChannel,discord.Thread], enable:bool, days:int, custom_message=None, for_user:discord.Member=None): + """Configurate productivity reminder.""" + command_name = 'productivity_config' + guild_id = ctx.guild.id + guild_data = get_guild_data(guild_id) + author = ctx.author + if is_feature_enabled('productivity',data=guild_data): + feature_data = guild_data['features']['productivity'] + if 'reminders' not in feature_data.keys(): + feature_data['reminders'] = dict() + if str(author.id) not in feature_data['reminders'].keys(): + feature_data['reminders'][str(author.id)] = dict() + if for_user: + if author.guild_permissions.administrator: + user_data = feature_data['reminders'][str(for_user.id)] + LOGGER.debug(f'{author.name} used /{command_name} {channel.mention} {for_user.name}.') + user_data['channel'] = channel.id + user_data['enabled'] = enable + user_data['cooldown'] = days#*60*60*24 + user_data['next_reminder'] = datetime.now().timestamp() + user_data['cooldown'] + user_data['custom_message'] = custom_message + save_guild_data(guild_id, guild_data) + # create a task for the user + if for_user.id in TASKS.keys(): + TASKS[str(for_user.id)].cancel() + TASKS[str(for_user.id)] = create_task(send_reminder(user_data['cooldown'],for_user,channel,user_data['custom_message'])) + LOGGER.debug(f'{author.name} set up productivity reminders for {for_user.name}.') + await ctx.send_response(f'Productivity reminders for {for_user.name} have been set to {channel.mention}.',ephemeral=True) + else: + LOGGER.debug(f'{author.name} tried to use /{command_name} {channel.mention} {for_user.name} but is not an administrator.') + await ctx.send_response('You must be an administrator to run the command with "for_user".',ephemeral=True) + else: + LOGGER.debug(f'{author.name} used /{command_name} {channel.mention}.') + user_data = feature_data['reminders'][str(author.id)] + user_data['channel'] = channel.id + user_data['enabled'] = enable + user_data['cooldown'] = days#*60*60*24 + user_data['next_reminder'] = datetime.now().timestamp() + user_data['cooldown'] + user_data['custom_message'] = custom_message + save_guild_data(guild_id, guild_data) + # create a task for the user + if str(author.id) in TASKS.keys(): + TASKS[str(author.id)].cancel() + TASKS[str(author.id)] = create_task(send_reminder(user_data['cooldown'],author,channel,user_data['custom_message'])) + LOGGER.debug(f'{author.name} set up productivity reminders for {author.name}.') + await ctx.send_response(f'Productivity reminders for {author.name} have been set to {channel.mention}.',ephemeral=True) + else: + LOGGER.debug(f'{author.name} tried to use /{command_name} {channel.mention} but the productivity feature is not enabled.') + await ctx.send_response('The productivity feature is not enabled.',ephemeral=True) + + @commands.Cog.listener() + async def on_message(self,message:discord.Message): + author = message.author + if author.id != self.bot.user.id: + guild_id = message.guild.id + guild_data = get_guild_data(guild_id) + if is_feature_enabled('productivity',data=guild_data): + channel = message.channel + feature_data = guild_data['features']['productivity'] + if 'reminders' in feature_data.keys(): + if str(author.id) in feature_data['reminders'].keys(): # If a user sat up its reminders + user_data = feature_data['reminders'][str(author.id)] + if user_data['enabled']: # If reminders are enabled for this user + reminder_channel_id = user_data['channel'] + if channel.id == reminder_channel_id: + user_data['next_reminder'] = datetime.now().timestamp() + user_data['cooldown'] + save_guild_data(guild_id, guild_data) + if author.id in TASKS.keys(): + TASKS[str(author.id)].cancel() + TASKS[str(author.id)] = create_task(send_reminder(user_data['cooldown'],author,channel,user_data['custom_message'])) + LOGGER.debug(f'{author.name} sent a message in {channel.name} and reminders are set up in {reminder_channel_id}.') + else: + LOGGER.debug(f'{author.name} sent a message in {channel.name} but reminders are set up in {reminder_channel_id}.') + + + \ No newline at end of file diff --git a/modules/setup/bot.py b/modules/setup/bot.py index 50ea39c..3a9c9d2 100644 --- a/modules/setup/bot.py +++ b/modules/setup/bot.py @@ -1,5 +1,7 @@ import discord from ..cogs.birthday import BirthdayCog +from ..cogs.productivity import ProductivityCog +from .logger import LOGGER bot_intents = discord.Intents.default() bot_intents.message_content = True @@ -10,4 +12,15 @@ bot_intents.reactions = True bot = discord.Bot(command_prefix='$', intents=bot_intents) -bot.add_cog(BirthdayCog(bot)) \ No newline at end of file +COGS = [BirthdayCog, ProductivityCog] + +for cog in COGS: + bot.add_cog(cog(bot)) + +def reload_feature(feature:str): + LOGGER.debug(bot.cogs) + LOGGER.debug(f'Reloading {feature}...') + bot.remove_cog(feature.capitalize()+'Cog') + cog = globals()[feature.capitalize()+'Cog'] + bot.add_cog(cog(bot)) + LOGGER.debug(f'{feature} reloaded!') diff --git a/modules/setup/data.py b/modules/setup/data.py index 797c9f1..03b5913 100644 --- a/modules/setup/data.py +++ b/modules/setup/data.py @@ -1,4 +1,7 @@ import os, json +from .logger import LOGGER + +FEATURES = ['birthday','productivity'] DATA_DIR = 'data/' @@ -12,4 +15,25 @@ def save_guild_data(guild_id:int, data): """Saves the dictionary of guild data to the guild's JSON data file.""" path = os.path.join(DATA_DIR, f'{guild_id}.json') with open(path, 'w') as f: - json.dump(data, f, indent=2) \ No newline at end of file + json.dump(data, f, indent=2) + LOGGER.debug(f'Saved data file for guild {guild_id}.') + +def create_guild_data(guild_id:int): + """Creates a new guild data file.""" + path = os.path.join(DATA_DIR, f'{guild_id}.json') + with open(path,'w') as f: + json.dump({'features': {key:{'enabled':True} for key in FEATURES}}, f, indent=2) + LOGGER.debug(f'Created data file for guild {guild_id}.') + +def delete_guild_data(guild_id:int): + """Deletes the guild data file.""" + path = os.path.join(DATA_DIR, f'{guild_id}.json') + os.remove(path) + LOGGER.debug(f'Deleted data file for guild {guild_id}.') + +def is_feature_enabled(feature:str,data=None,guild_id=0): + """Returns True if the feature is enabled in the guild, False otherwise.""" + if data: + return data['features'][feature]['enabled'] + else: + return get_guild_data(guild_id)['features'][feature]['enabled'] \ No newline at end of file diff --git a/modules/setup/logger.py b/modules/setup/logger.py index 4055e58..b5bc5cb 100644 --- a/modules/setup/logger.py +++ b/modules/setup/logger.py @@ -3,6 +3,6 @@ import logging logging.basicConfig(level=logging.INFO, format='[%(levelname)s] %(asctime)s: %(name)s: %(message)s') LOGGER = logging.getLogger('botmafieux') LOGGER.setLevel(logging.DEBUG) -handler = logging.FileHandler(filename='botmafieux.log', encoding='utf-8', mode='w') +handler = logging.FileHandler(filename='botmafieux.log', encoding='utf-8', mode='a') handler.setFormatter(logging.Formatter('%(asctime)s:%(levelname)s:%(name)s: %(message)s')) LOGGER.addHandler(handler) \ No newline at end of file