Merge pull request #59 from StudyLions/rewrite

src/tracking/voice/cog.py

@@ -1,5 +1,6 @@
 from typing import Optional
 import asyncio
+import itertools
 import datetime as dt
 
 import discord
@@ -71,7 +72,7 @@ class VoiceTrackerCog(LionCog):
             locked=0,
             actual=0,
             channels=0,
-            cached=len(VoiceSession._sessions_),
+            cached=sum(len(gsessions) for gsessions in VoiceSession._sessions_.values()),
             initial_event=self.initialised,
             lock=self.tracking_lock
         )
@@ -92,7 +93,7 @@ class VoiceTrackerCog(LionCog):
         data['channels'] = len(channels)
 
         for guild in self.bot.guilds:
-            for channel in guild.voice_channels:
+            for channel in itertools.chain(guild.voice_channels, guild.stage_channels):
                 if not self.is_untracked(channel):
                     for member in channel.members:
                         if member.voice and not member.bot:
@@ -143,7 +144,7 @@ class VoiceTrackerCog(LionCog):
         ...
 
     # ----- Cog API -----
-    def get_session(self, guildid, userid, **kwargs) -> Optional[VoiceSession]:
+    def get_session(self, guildid, userid, **kwargs):
         """
         Get the VoiceSession for the given member.
 
@@ -325,6 +326,8 @@ class VoiceTrackerCog(LionCog):
             active = self.active_sessions.pop(guild.id, {}).values()
             for session in active:
                 session.cancel()
+            # Clear registry
+            VoiceSession._sessions_.pop(guild.id, None)
 
             # Update untracked channel information for this guild
             self.untracked_channels.pop(guild.id, None)
@@ -332,7 +335,7 @@ class VoiceTrackerCog(LionCog):
 
             # Read tracked voice states
             states = {}
-            for channel in guild.voice_channels:
+            for channel in itertools.chain(guild.voice_channels, guild.stage_channels):
                 if not self.is_untracked(channel):
                     for member in channel.members:
                         if member.voice and not member.bot:
@@ -390,7 +393,7 @@ class VoiceTrackerCog(LionCog):
             # Read and save the tracked voice states of all visible voice channels
             states = {}
             for guild in self.bot.guilds:
-                for channel in guild.voice_channels:
+                for channel in itertools.chain(guild.voice_channels, guild.stage_channels):
                     if not self.is_untracked(channel):
                         for member in channel.members:
                             if member.voice and not member.bot:

src/tracking/voice/session.py

@@ -253,6 +253,8 @@ class VoiceSession:
             f"Expiring voice session for member <uid:{self.userid}> in guild <gid:{self.guildid}> "
             f"and channel <cid:{self.state.channelid}>."
         )
+        # TODO: Would be better not to close the session and wipe the state
+        # Instead start a new PENDING session.
         await self.close()
 
     async def update(self, new_state: Optional[TrackedVoiceState] = None, new_rate: Optional[int] = None):

src/tracking/voice/settings.py

@@ -457,6 +457,9 @@ class VoiceTrackerConfigUI(ConfigUI):
     @select(
         cls=ChannelSelect,
         placeholder="UNTRACKED_CHANNELS_PLACEHOLDER",
+        channel_types=[
+            discord.enums.ChannelType.voice, discord.enums.ChannelType.stage_voice, discord.enums.ChannelType.category
+        ],
         min_values=0, max_values=25
     )
     async def untracked_channels_menu(self, selection: discord.Interaction, selected):