Merge pull request #314 from PyBotDevs/user-data-library

Migrate UserData system completely to a framework module

Author: notsniped

cogs/economy.py

@@ -2,14 +2,13 @@
 
 # Imports
 import discord
-import json
 import random
 import math
 import utils.logger
 import asyncio
 import framework.isobot.currency
 from framework.isobot.shop import ShopData
-from framework.isobot.db import levelling, items
+from framework.isobot.db import levelling, items, userdata
 from random import randint
 from discord import option, ApplicationContext
 from discord.ext import commands
@@ -19,6 +18,7 @@
 currency = framework.isobot.currency.CurrencyAPI("database/currency.json", "logs/currency.log")
 levelling = levelling.Levelling()
 items = items.Items()
+userdata = userdata.UserData()
 shop_data = ShopData("config/shop.json")
 all_item_ids = shop_data.get_item_ids()
 shopitem = shop_data.get_raw_data()
@@ -32,19 +32,6 @@
     "Doctor"
 ]
 
-with open("database/user_data.json", 'r') as f: userdat = json.load(f)
-
-def save():
-    with open("database/user_data.json", 'w+') as f: json.dump(userdat, f, indent=4)
-
-# Functions
-def new_userdat(id: int):
-    if str(id) not in userdat.keys(): 
-        userdat[str(id)] = {"work_job": None}
-        save()
-        return 0
-    else: return 1
-
 # Commands
 class Economy(commands.Cog):
     def __init__(self, bot):
@@ -410,16 +397,17 @@ async def give(self, ctx: ApplicationContext, user:discord.User, amount:int):
     )
     @commands.cooldown(1, 1800, commands.BucketType.user)
     async def work(self, ctx: ApplicationContext):
-        if userdat[str(ctx.author.id)]["work_job"] == None: return await ctx.respond("You don't currently have a job! Join one by using the `/work_select` command.", ephemeral=True)
-        if userdat[str(ctx.author.id)]["work_job"] == "Discord mod": i = randint(5000, 10000)
-        elif userdat[str(ctx.author.id)]["work_job"] == "YouTuber": i = randint(10000, 15000)
-        elif userdat[str(ctx.author.id)]["work_job"] == "Streamer": i = randint(12000, 18000)
-        elif userdat[str(ctx.author.id)]["work_job"] == "Developer": i = randint(20000, 40000)
-        elif userdat[str(ctx.author.id)]["work_job"] == "Scientist": i = randint(50000, 100000)
-        elif userdat[str(ctx.author.id)]["work_job"] == "Engineer": i = randint(100000, 175000)
-        elif userdat[str(ctx.author.id)]["work_job"] == "Doctor": i = randint(200000, 300000)
+        job_name = userdata.fetch(ctx.author.id, "work_job")
+        if job_name == None: return await ctx.respond("You don't currently have a job! Join one by using the `/work_select` command.", ephemeral=True)
+        if job_name == "Discord mod": i = randint(5000, 10000)
+        elif job_name == "YouTuber": i = randint(10000, 15000)
+        elif job_name == "Streamer": i = randint(12000, 18000)
+        elif job_name == "Developer": i = randint(20000, 40000)
+        elif job_name == "Scientist": i = randint(50000, 100000)
+        elif job_name == "Engineer": i = randint(100000, 175000)
+        elif job_name == "Doctor": i = randint(200000, 300000)
         currency.add(ctx.author.id, i)
-        await ctx.respond(f'{ctx.author.mention} worked for a 30-minute shift as a {userdat[str(ctx.author.id)]["work_job"]} and earned {i} coins.')
+        await ctx.respond(f'{ctx.author.mention} worked for a 30-minute shift as a {job_name} and earned {i} coins.')
 
     @commands.slash_command(
         name="work_list",
@@ -447,8 +435,7 @@ async def work_select(self, ctx: ApplicationContext, job: str):
         elif job == "Scientist" and levelling.get_level(ctx.author.id) < 20: return await ctx.respond("You currently do not have the required level to perform this job!", ephemeral=True)
         elif job == "Engineer" and levelling.get_level(ctx.author.id) < 25: return await ctx.respond("You currently do not have the required level to perform this job!", ephemeral=True)
         elif job == "Doctor" and levelling.get_level(ctx.author.id) < 40: return await ctx.respond("You currently do not have the required level to perform this job!", ephemeral=True)
-        userdat[str(ctx.author.id)]["work_job"] = job
-        save()
+        userdata.set(ctx.author.id, "work_job", job)
         localembed = discord.Embed(title="New job!", description=f"You are now working as a {job}!")
         await ctx.respond(embed=localembed)
 
@@ -457,9 +444,8 @@ async def work_select(self, ctx: ApplicationContext, job: str):
         description="Quit your job."
     )
     async def work_resign(self, ctx: ApplicationContext):
-        if userdat[str(ctx.author.id)]["work_job"] is None: return await ctx.respond("You can't quit your job if you don't already have one!", ephemeral=True)
-        userdat[str(ctx.author.id)]["work_job"] = None
-        save()
+        if userdata.fetch(ctx.author.id, "work_job") is None: return await ctx.respond("You can't quit your job if you don't already have one!", ephemeral=True)
+        userdata.set(ctx.author.id, "work_job", None)
         localembed = discord.Embed(title="Resignation", description="You have successfully resigned from your job.")
         await ctx.respond(embed=localembed)
 

framework/isobot/db/userdata.py

@@ -0,0 +1,42 @@
+"""The framework module library used for managing general user data."""
+
+# Imports
+import json
+
+# Functions
+class UserData():
+    """Used to initialize the UserData system."""
+    def __init__(self):
+        print("[framework/db/UserData] UserData library initialized.")
+    
+    def load(self) -> dict:
+        """Fetches and returns the latest data from the levelling database."""  
+        with open("database/user_data.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/user_data.json", 'w+', encoding="utf8") as f: json.dump(data, f, indent=4)
+        return 0
+    
+    def generate(self, user_id: int) -> int:
+        """Generates a new data key for the specified user.\n
+        Returns `0` if the request was successful, returns `1` if the data key already exists."""
+        userdat = self.load()
+        if str(user_id) not in userdat.keys(): 
+            userdat[str(user_id)] = {"work_job": None}
+            self.save(userdat)
+            return 0
+        else: return 1
+
+    def fetch(self, user_id: int, key: str) -> str:
+        """Fetches the vakue of a data key, from a specific user."""
+        userdat = self.load()
+        return userdat[str(user_id)][key]
+    
+    def set(self, user_id: int, key: str, value) -> int:
+        """Sets a new value for a data key, for a specific user."""
+        userdat = self.load()
+        userdat[str(user_id)][key] = value
+        self.save(userdat)
+        return 0

main.py

@@ -11,10 +11,9 @@
 from math import floor
 from random import randint
 from framework.isobot import currency, colors, settings
-from framework.isobot.db import levelling, items
+from framework.isobot.db import levelling, items, userdata
 from discord import ApplicationContext, option
 from discord.ext import commands
-from cogs.economy import new_userdat
 from cogs.isocoin import create_isocoin_key
 
 # Slash option types:
@@ -57,6 +56,7 @@ def save():
 settings = settings.Configurator()
 levelling = levelling.Levelling()
 items = items.Items()
+userdata = userdata.UserData()
 
 # Theme Loader
 themes = False  # True: enables themes; False: disables themes;
@@ -96,7 +96,7 @@ async def on_message(ctx):
     currency.new_wallet(ctx.author.id)
     currency.new_bank(ctx.author.id)
     create_isocoin_key(ctx.author.id)
-    new_userdat(ctx.author.id)
+    userdata.generate(ctx.author.id)
     settings.generate(ctx.author.id)
     items.generate(ctx.author.id)
     levelling.generate(ctx.author.id)