Restrictions: Significantly clean up code

cogs/logs.py

@@ -1,6 +1,7 @@
 import discord
 import json
 import config
+from helpers.restrictions import get_user_restrictions
 
 
 class Logs:
@@ -43,14 +44,10 @@ class Logs:
               f"🕓 Account age: {age}\n"\
               f"🏷 __User ID__: {member.id}"
 
-        # Taken from kurisu source.
+        # Handles user restrictions
-        # Blame ihaveamac, not me.
+        # Basically, gives back muted role to users that leave with it.
-        with open("data/restrictions.json", "r") as f:
+        rsts = get_user_restrictions(member.id)
-            rsts = json.load(f)
+        roles = [discord.utils.get(member.guild.roles, id=rst) for rst in rsts]
-        if str(member.id) in rsts:
-            roles = []
-            for rst in rsts[str(member.id)]:
-                roles.append(discord.utils.get(member.guild.roles, id=rst))
         await member.add_roles(*roles)
 
         # Real hell zone.

cogs/mod.py

@@ -57,7 +57,7 @@ class Mod:
         log_channel = self.bot.get_channel(config.log_channel)
         await log_channel.send(chan_message)
         await ctx.send(f"{target.mention} can no longer speak.")
-        await add_restriction(target, config.mute_role)
+        add_restriction(target.id, config.mute_role)
 
     @commands.guild_only()
     @commands.check(check_if_staff)
@@ -76,7 +76,7 @@ class Mod:
         log_channel = self.bot.get_channel(config.log_channel)
         await log_channel.send(chan_message)
         await ctx.send(f"{target.mention} can now speak again.")
-        await remove_restriction(target, config.mute_role)
+        remove_restriction(target.id, config.mute_role)
 
     @commands.guild_only()
     @commands.bot_has_permissions(kick_members=True)

helpers/restrictions.py

@@ -1,25 +1,42 @@
 import json
 
 
-def add_restriction(self, member, rst):
+def get_restrictions():
-    # from kurisu source, credits go to ihaveamac
     with open("data/restrictions.json", "r") as f:
-        rsts = json.load(f)
+        return json.load(f)
-    if str(member.id) not in rsts:
-        rsts[str(member.id)] = []
-    if rst not in rsts[str(member.id)]:
-        rsts[str(member.id)].append(rst)
-    with open("data/restrictions.json", "w") as f:
-        json.dump(rsts, f)
 
 
-def remove_restriction(self, member, rst):
+def set_restrictions(contents):
-    # from kurisu source, credits go to ihaveamac
+    with open("data/restrictions.json", "w") as f:
+        f.write(contents)
+
+
+def get_user_restrictions(uid):
+    uid = str(uid)
     with open("data/restrictions.json", "r") as f:
         rsts = json.load(f)
-    if str(member.id) not in rsts:
+        if uid in rsts:
-        rsts[str(member.id)] = []
+            return rsts[uid]
-    if rst in rsts[str(member.id)]:
+        return []
-        rsts[str(member.id)].remove(rst)
+
-    with open("data/restrictions.json", "w") as f:
+
-        json.dump(rsts, f)
+def add_restriction(uid, rst):
+    # mostly from kurisu source, credits go to ihaveamac
+    uid = str(uid)
+    rsts = get_restrictions()
+    if uid not in rsts:
+        rsts[uid] = []
+    if rst not in rsts[uid]:
+        rsts[uid].append(rst)
+    set_restrictions(json.dumps(rsts))
+
+
+def remove_restriction(uid, rst):
+    # mostly from kurisu source, credits go to ihaveamac
+    uid = str(uid)
+    rsts = get_restrictions()
+    if uid not in rsts:
+        rsts[uid] = []
+    if rst in rsts[uid]:
+        rsts[uid].remove(rst)
+    set_restrictions(json.dumps(rsts))