Switch to custom group for approved use and added data migration

lseen/__init__.py

@@ -1,5 +1,7 @@
 from .lseen import LastSeen
 
 
-def setup(bot):
-    bot.add_cog(LastSeen(bot))
+async def setup(bot):
+    seen = LastSeen(bot)
+    bot.add_cog(seen)
+    await seen.initialize()

lseen/lseen.py

@@ -1,3 +1,5 @@
+import asyncio
+import logging
 from datetime import datetime
 from typing import Literal
 
@@ -8,6 +10,8 @@ from redbot.core.bot import Red
 from redbot.core.commands import Cog
 from redbot.core.utils import AsyncIter
 
+log = logging.getLogger("red.fox_v3.lseen")
+
 
 class LastSeen(Cog):
     """
@@ -22,13 +26,15 @@ class LastSeen(Cog):
         super().__init__()
         self.bot = bot
         self.config = Config.get_conf(self, identifier=9811198108111121, force_registration=True)
-        default_global = {}
-        default_guild = {"enabled": False}
+        default_global = {"migrated_v2": False}
+        default_guild = {"enabled": None}
         default_member = {"seen": None}
 
         self.config.register_global(**default_global)
         self.config.register_guild(**default_guild)
-        self.config.register_member(**default_member)
+        # self.config.register_member(**default_member)
+        self.config.init_custom("CustomMember", 2)
+        self.config.register_custom("CustomMember", **default_member)
 
     async def red_delete_data_for_user(
         self,
@@ -41,7 +47,20 @@ class LastSeen(Cog):
 
         async for guild_id, guild_data in AsyncIter(all_members.items(), steps=100):
             if user_id in guild_data:
-                await self.config.member_from_ids(guild_id, user_id).clear()
+                await self.config.custom("CustomMember", guild_id, str(user_id)).clear()
+
+    async def _initalize_tracking(self, guild: discord.Guild):
+        now = datetime.utcnow().isoformat()
+        online_members = {
+            member.id: {"seen": now}
+            for member in guild.members
+            if member.status != self.offline_status
+        }
+        async with self.config.custom("CustomMember", guild.id).all() as cm:
+            cm.update(online_members)
+        # for member in guild.members:
+        #     if member.status != self.offline_status:
+        #         await self.config.custom("CustomMember", guild.id, member.id).seen.set(now)
 
     @staticmethod
     def get_date_time(s):
@@ -57,9 +76,19 @@ class LastSeen(Cog):
     @lset.command(name="toggle")
     async def lset_toggle(self, ctx: commands.Context):
         """Toggles tracking seen for this server"""
-        enabled = not await self.config.guild(ctx.guild).enabled()
+        enabled = await self.config.guild(ctx.guild).enabled()
+        if enabled is None:
+            needs_init = True
+            enabled = True
+        else:
+            needs_init = False
+            enabled = not enabled
+
         await self.config.guild(ctx.guild).enabled.set(enabled)
 
+        if needs_init:
+            asyncio.create_task(self._initalize_tracking(ctx.guild))
+
         await ctx.maybe_send_embed(
             "Seen for this server is now {}".format("Enabled" if enabled else "Disabled")
         )
@@ -72,8 +101,12 @@ class LastSeen(Cog):
 
         if member.status != self.offline_status:
             last_seen = datetime.utcnow()
+            if await self.config.guild(ctx.guild).enabled():
+                await self.config.custom("CustomMember", member.guild.id, member.id).seen.set(
+                    last_seen.isoformat()
+                )
         else:
-            last_seen = await self.config.member(member).seen()
+            last_seen = await self.config.custom("CustomMember", member.guild.id, member.id).seen()
             if last_seen is None:
                 await ctx.maybe_send_embed("I've never seen this user")
                 return
@@ -93,4 +126,20 @@ class LastSeen(Cog):
                 return
             if not await self.config.guild(after.guild).enabled():
                 return
-            await self.config.member(before).seen.set(datetime.utcnow().isoformat())
+            await self.config.custom("CustomMember", before.guild.id, before.id).seen.set(
+                datetime.utcnow().isoformat()
+            )
+
+    async def initialize(self):
+        if not await self.config.migrated_v2():
+            all_members = await self.config.all_members()
+            async for guild_id, guild_data in AsyncIter(all_members.items(), steps=100):
+                async with self.config.custom("CustomMember").all() as cm:
+                    cm.update(
+                        {
+                            str(guild_id): {str(member_id): {"seen": member_data["seen"]}}
+                            for member_id, member_data in guild_data.items()
+                        }
+                    )
+            log.info("LastSeen migrated to V2")
+            await self.config.migrated_v2.set(True)