Move all AFK and presence database functions to a new framework module

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))

cogs/utils.py

@@ -10,14 +10,15 @@
 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()
 currency = currency.CurrencyAPI("database/currency.json", "logs/currency.log")
 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']} - <t:{math.floor(presence['time'])}>"
         localembed = discord.Embed(

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

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,15 +25,13 @@
 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()
 
 #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 <t:{floor(presence[str(ctx.guild.id)][str(user)]['time'])}:R>: {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()