diff --git a/lseen/__init__.py b/lseen/__init__.py index 716f34d..53df35e 100644 --- a/lseen/__init__.py +++ b/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() diff --git a/lseen/lseen.py b/lseen/lseen.py index f77fbca..b9e1727 100644 --- a/lseen/lseen.py +++ b/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): @@ -55,9 +74,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") ) @@ -70,8 +99,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)