diff --git a/cogs/afk.py b/cogs/afk.py index 80279b5f..c895dd92 100644 --- a/cogs/afk.py +++ b/cogs/afk.py @@ -3,30 +3,16 @@ # Imports import discord import json -import time from discord import option, ApplicationContext, SlashCommandGroup from discord.ext import commands +from framework.isobot.db.presence import Presence # Variables with open("database/presence.json", 'r', encoding="utf-8") as f: presence = json.load(f) - -def save(): - with open("database/presence.json", 'w+', encoding="utf-8") as f: json.dump(presence, f) - -# Functions -def get_presence(user_id: int, guild_id: int) -> dict: - """Returns a `dict` of the specified user's current AFK status in the guild.""" - if str(user_id) in presence[str(guild_id)]: - return { - "afk": True, - "response": presence[str(guild_id)][str(user_id)]['response'], - "time": presence[str(guild_id)][str(user_id)]['time'] - } - else: - return False +presence = Presence() # Commands -class Presence(commands.Cog): +class PresenceCog(commands.Cog): def __init__(self, bot): self.bot = bot @@ -37,11 +23,8 @@ def __init__(self, bot): description="Sets your AFK status with a custom response" ) @option(name="response", description="What do you want your AFK response to be?", type=str, default="I'm AFK") - async def afk_set(self, ctx: ApplicationContext, response:str="I'm AFK"): - exctime = time.time() - if str(ctx.guild.id) not in presence: presence[str(ctx.guild.id)] = {} - presence[str(ctx.guild.id)][str(ctx.author.id)] = {"type": "afk", "time": exctime, "response": response} - save() + async def afk_set(self, ctx: ApplicationContext, response: str="I'm AFK"): + presence.add_afk(ctx.guild.id, ctx.user.id, response) localembed = discord.Embed(title=f"{ctx.author.display_name} is now AFK.", description=f"Response: {response}", color=discord.Color.dark_orange()) await ctx.respond(embed=localembed) @@ -50,11 +33,9 @@ async def afk_set(self, ctx: ApplicationContext, response:str="I'm AFK"): description="Removes your AFK status" ) async def afk_remove(self, ctx: ApplicationContext): - try: - del presence[str(ctx.guild.id)][str(ctx.author.id)] - save() - await ctx.respond(f"Alright {ctx.author.mention}, I've removed your AFK.") - except KeyError: return await ctx.respond("You weren't previously AFK.", ephemeral=True) + status = presence.remove_afk(ctx.guild.id, ctx.author.id) + if status == 0: return await ctx.respond(f"Alright {ctx.author.mention}, I've removed your AFK.") + elif status == 1: return await ctx.respond("You weren't previously AFK.", ephemeral=True) @afk_system.command( name="mod_remove", @@ -63,11 +44,9 @@ async def afk_remove(self, ctx: ApplicationContext): @option(name="user", description="Whose AFK status do you want to remove?", type=discord.User) async def afk_mod_remove(self, ctx: ApplicationContext, user:discord.User): if not ctx.author.guild_permissions.manage_messages: return await ctx.respond("You don't have the required permissions to use this.", ephemeral=True) - try: - del presence[str(ctx.guild.id)][str(user.id)] - save() - await ctx.respond(f"{user.display_name}'s AFK has been removed.") - except KeyError: return await ctx.respond("That user isn't AFK.", ephemeral=True) + status = presence.remove_afk(ctx.guild.id, user.id) + if status == 0: return await ctx.respond(f"{user.display_name}'s AFK has been removed.") + elif status == 1: return await ctx.respond("That user isn't AFK.", ephemeral=True) # Cog Initialization -def setup(bot): bot.add_cog(Presence(bot)) +def setup(bot): bot.add_cog(PresenceCog(bot)) diff --git a/cogs/utils.py b/cogs/utils.py index 7dd41baa..721f40b4 100644 --- a/cogs/utils.py +++ b/cogs/utils.py @@ -10,7 +10,7 @@ from framework.isobot.db import levelling from discord import option, ApplicationContext from discord.ext import commands -from cogs.afk import get_presence +from framework.isobot.db.presence import Presence # Variables color = discord.Color.random() @@ -18,6 +18,7 @@ levelling = levelling.Levelling() openai.api_key = os.getenv("chatgpt_API_KEY") chatgpt_conversation = dict() +_presence = Presence() # Commands class Utils(commands.Cog): @@ -91,7 +92,7 @@ async def whoami(self, ctx: ApplicationContext, user: discord.User=None): registered = user.joined_at.strftime("%b %d, %Y, %T") pfp = user.avatar localembed_desc = f"`AKA` {displayname}" - presence = get_presence(ctx.author.id, ctx.guild.id) + presence = _presence.get_presence(ctx.guild.id, ctx.user.id) if presence is not False: localembed_desc += f"\n`🌙 AFK` {presence['response']} - " localembed = discord.Embed( diff --git a/framework/isobot/db/presence.py b/framework/isobot/db/presence.py new file mode 100644 index 00000000..218e236a --- /dev/null +++ b/framework/isobot/db/presence.py @@ -0,0 +1,55 @@ +"""The framework module library used for managing user presence and AFK data.""" + +# Imports +import json +import time + +# Functions +class Presence(): + """Used to initialize the User Presence system.""" + def __init__(self): + print("[framework/db/UserData] Presence database initialized.") + + def load(self) -> dict: + """Fetches and returns the latest data from the presence database.""" + with open("database/presence.json", 'r', encoding="utf8") as f: db = json.load(f) + return db + + def save(self, data: dict) -> int: + """Dumps all cached data to your local machine.""" + with open("database/presence.json", 'w+', encoding="utf8") as f: json.dump(data, f) + return 0 + + def add_afk(self, guild_id: int, user_id: int, response: str) -> int: + """Generates a new afk data request in the guild for the specified user.""" + presence = self.load() + exctime = time.time() + if str(guild_id) not in presence: presence[str(guild_id)] = {} + presence[str(guild_id)][str(user_id)] = {"type": "afk", "time": exctime, "response": response} + self.save(presence) + return 0 + + def remove_afk(self, guild_id: int, user_id: int) -> int: + """Clears the specified user's AFK data.\n\nReturns `0` if the request was successful, returns 1 if the user's AFK is not present in the database (equivalent to `KeyError`).""" + presence = self.load() + try: + del presence[str(guild_id)][str(user_id)] + self.save(presence) + return 0 + except KeyError: return 1 + + def get_presence(self, guild_id: int, user_id: int) -> dict: + """Returns a `dict` of the specified user's current AFK status in the guild. Returns `1` if the user is not in the presence database.""" + presence = self.load() + if str(user_id) in presence[str(guild_id)]: + return { + "afk": True, + "response": presence[str(guild_id)][str(user_id)]['response'], + "time": presence[str(guild_id)][str(user_id)]['time'] + } + else: return 1 + + def get_raw(self) -> dict: + """Returns the entire raw presence database. Can be useful for complex data operations on the database.""" + presence = self.load() + return presence diff --git a/main.py b/main.py index 32b26f7d..419560fc 100644 --- a/main.py +++ b/main.py @@ -11,7 +11,7 @@ from math import floor from random import randint from framework.isobot import currency, colors, settings -from framework.isobot.db import levelling, items, userdata, automod +from framework.isobot.db import levelling, items, userdata, automod, presence as _presence from discord import ApplicationContext, option from discord.ext import commands from cogs.isocoin import create_isocoin_key @@ -25,7 +25,6 @@ client = discord.Bot() color = discord.Color.random() with open('config/shop.json', 'r', encoding="utf-8") as f: shopitem = json.load(f) -with open('database/presence.json', 'r', encoding="utf-8") as f: presence = json.load(f) with open('config/commands.json', 'r', encoding="utf-8") as f: commandsdb = json.load(f) with open('database/automod.json', 'r', encoding="utf-8") as f: automod_config = json.load(f) cmd_list = commandsdb.keys() @@ -33,7 +32,6 @@ #Pre-Initialization Commands def save(): """Dumps all cached data to the local databases.""" - with open('database/presence.json', 'w+', encoding="utf-8") as e: json.dump(presence, e, indent=4) with open('database/automod.json', 'w+', encoding="utf-8") as e: json.dump(automod_config, e, indent=4) if not os.path.isdir("logs"): @@ -58,6 +56,7 @@ def save(): items = items.Items() userdata = userdata.UserData() automod = automod.Automod() +_presence = _presence.Presence() # Theme Loader themes = False # True: enables themes; False: disables themes; @@ -103,6 +102,7 @@ async def on_message(ctx): levelling.generate(ctx.author.id) automod.generate(ctx.guild.id) uList = list() + presence = _presence.get_raw() if str(ctx.guild.id) in presence: for userid in presence[str(ctx.guild.id)].keys(): uList.append(userid) else: pass @@ -111,8 +111,7 @@ async def on_message(ctx): fetch_user = client.get_user(id(user)) await ctx.channel.send(f"{fetch_user.display_name} went AFK : {presence[str(ctx.guild.id)][str(user)]['response']}") if str(ctx.guild.id) in presence and str(ctx.author.id) in presence[str(ctx.guild.id)]: - del presence[str(ctx.guild.id)][str(ctx.author.id)] - save() + _presence.remove_afk(ctx.guild.id, ctx.author.id) m1 = await ctx.channel.send(f"Welcome back {ctx.author.mention}. Your AFK has been removed.") await asyncio.sleep(5) await m1.delete()