Add role persistence cog (#23)

robocop_ng/__main__.py

@@ -46,6 +46,7 @@ wanted_jsons = [
     "data/robocronptab.json",
     "data/userlog.json",
     "data/invites.json",
+    "data/persistent_roles.json"
 ]
 
 intents = discord.Intents.all()

robocop_ng/cogs/role_persistence.py

@@ -0,0 +1,29 @@
+from discord import RawMemberRemoveEvent, Member
+from discord.ext.commands import Cog
+
+from robocop_ng.helpers.roles import add_user_roles, get_user_roles
+
+
+class RolePersistence(Cog):
+    def __init__(self, bot):
+        self.bot = bot
+
+    @Cog.listener
+    async def on_raw_member_remove(self, payload: RawMemberRemoveEvent):
+        save_roles = []
+        for role in payload.user.roles:
+            if role.is_assignable() and \
+                    not role.is_default() and not role.is_premium_subscriber() and \
+                    not role.is_bot_managed() and not role.is_integration():
+                save_roles.append(role.id)
+
+        if len(save_roles) > 0:
+            add_user_roles(payload.user.id, save_roles)
+
+    @Cog.listener
+    async def on_member_join(self, member: Member):
+        user_roles = get_user_roles(member.id)
+        if len(user_roles) > 0:
+            user_roles = [member.guild.get_role(int(role)) for role in user_roles
+                          if member.guild.get_role(int(role)) is not None]
+            await member.add_roles(*user_roles, reason="Restoring old roles from `role_persistence`.")

robocop_ng/helpers/roles.py

@@ -0,0 +1,34 @@
+import json
+import os.path
+
+PERSISTENT_ROLES_FILE = "data/persistent_roles.json"
+
+
+def get_persistent_roles() -> dict[str, list[str]]:
+    if os.path.isfile(PERSISTENT_ROLES_FILE):
+        with open(PERSISTENT_ROLES_FILE, "r") as f:
+            return json.load(f)
+    return {}
+
+
+def set_persistent_roles(contents: dict[str, list[str]]):
+    with open(PERSISTENT_ROLES_FILE, "w") as f:
+        json.dump(contents, f)
+
+
+def add_user_roles(uid: int, roles: list[int]):
+    uid = str(uid)
+    roles = [str(x) for x in roles]
+
+    persistent_roles = get_persistent_roles()
+    persistent_roles[uid] = roles
+    set_persistent_roles(persistent_roles)
+
+
+def get_user_roles(uid: int) -> list[str]:
+    uid = str(uid)
+    with open(PERSISTENT_ROLES_FILE, "r") as f:
+        roles = json.load(f)
+        if uid in roles:
+            return roles[uid]
+        return []