cleaned up edge.py

sizebot/cogs/edge.py

@@ -1,14 +1,15 @@
 # The "edge" cog allows and  guild owners to set a largest/smallest user for their server.
 # It does this by seeing if they are the largest or smallest user in the guild, and if they aren't setting their height
 # to either 1.1x or 0.9x of the largest or smallest user in the guild, respectively.
-from typing import Any
+from collections.abc import Iterable
 
 import logging
-from sizebot.lib.errors import GuildNotFoundException, UserNotFoundException
+from dataclasses import dataclass
 
 import discord
 from discord.ext import commands
 
+from sizebot.lib.errors import GuildNotFoundException, UserNotFoundException
 from sizebot.lib import guilddb, userdb, nickmanager
 from sizebot.lib.checks import is_mod
 from sizebot.lib.digidecimal import Decimal
@@ -18,36 +19,47 @@
 logger = logging.getLogger("sizebot")
 
 
-# TODO: CamelCase
-def getUserSizes(g: discord.Guild) -> Any:
-    # Find the largest and smallest current users.
-    # TODO: Can this be faster?
-    # Like, if you sorted the users by size, would that make this faster?
-    # It's 1:30AM and I don't want to check.
-    smallestuser = None
-    smallestsize = SV.infinity
-    largestuser = None
-    largestsize = SV(0)
-    allusers = {}
-    for _, userid in userdb.list_users(guildid=g.id):
-        member = g.get_member(userid)
-        if not (member and str(member.status) != "offline"):
+@dataclass
+class UserSize:
+    id: int | None
+    size: SV
+
+
+@dataclass
+class UserSizes:
+    smallest: UserSize
+    largest: UserSize
+    users: list[UserSize]
+
+
+def _get_active_users(guild: discord.Guild) -> Iterable[UserSize]:
+    for guildid, userid in userdb.list_users(guildid=guild.id):
+        member = guild.get_member(userid)
+        if member is None or str(member.status) == "offline":
             continue
-        userdata = userdb.load(g.id, userid)
+        userdata = userdb.load(guildid, userid)
         if userdata.height == 0 or userdata.height == SV.infinity:
             continue
         if not userdata.is_active:
             continue
-        if userdata.height > largestsize:
-            largestuser = userid
-            largestsize = userdata.height
-        if userdata.height < smallestsize:
-            smallestuser = userid
-            smallestsize = userdata.height
-        allusers[userid] = userdata.height
-    return {"smallest": {"id": smallestuser, "size": smallestsize},
-            "largest": {"id": largestuser, "size": largestsize},
-            "users": allusers}
+        yield UserSize(userid, userdata.height)
+
+
+def _get_user_sizes(g: discord.Guild) -> UserSizes:
+    # Find the largest and smallest current users.
+    # TODO: Can this be faster?
+    # Like, if you sorted the users by size, would that make this faster?
+    # It's 1:30AM and I don't want to check.
+    smallest = UserSize(None, SV.infinity)
+    largest = UserSize(None, SV(0))
+    users: list[UserSize] = []
+    for user in _get_active_users(g):
+        if user.size > largest.size:
+            largest = user
+        if user.size < smallest.size:
+            smallest = user
+        users.append(user)
+    return UserSizes(smallest, largest, users)
 
 
 class EdgeCog(commands.Cog):
@@ -131,19 +143,19 @@ async def clearlargest(self, ctx: BotContext):
     @is_mod()
     async def edgedebug(self, ctx: BotContext):
         userdata = userdb.load(ctx.guild.id, ctx.author.id)
-        usersizes = getUserSizes(ctx.guild)
+        usersizes = _get_user_sizes(ctx.guild)
         guilddata = guilddb.load(ctx.guild.id)
         sm = guilddata.small_edge
         lg = guilddata.large_edge
 
         outstring = f"**CURRENT USER:**\nID: `{ctx.author.id}`\nHeight: `{userdata.height}`\n\n"
         outstring += f"**EDGES:**\nSmallest: {sm}\nLargest: {lg}\n\n"
-        outstring += f"**SMALLEST USER:**\nID: `{usersizes['smallest']['id']}`\nHeight: `{usersizes['smallest']['size']}`\n\n"
-        outstring += f"**LARGEST USER:**\nID: `{usersizes['largest']['id']}`\nHeight: `{usersizes['largest']['size']}`\n\n"
+        outstring += f"**SMALLEST USER:**\nID: `{usersizes.smallest.id}`\nHeight: `{usersizes.smallest.size}`\n\n"
+        outstring += f"**LARGEST USER:**\nID: `{usersizes.largest.id}`\nHeight: `{usersizes.largest.size}`\n\n"
         outstring += "**ALL USERS:**\n"
 
-        for pair in usersizes['users'].items():
-            outstring += f"`{pair[0]}`: {pair[1]}\n"
+        for u in usersizes.users:
+            outstring += f"`{u.id}`: {u.size}\n"
 
         await ctx.send(outstring)
 
@@ -168,29 +180,23 @@ async def on_message(self, m: discord.Message):
         except UserNotFoundException:
             return
 
-        usersizes = getUserSizes(m.guild)
-        smallestuser = usersizes["smallest"]["id"]
-        smallestsize = usersizes["smallest"]["size"]
-        largestuser = usersizes["largest"]["id"]
-        largestsize = usersizes["largest"]["size"]
+        usersizes = _get_user_sizes(m.guild)
 
         if sm == m.author.id:
-            if m.author.id == smallestuser:
+            if m.author.id == usersizes.smallest.id:
                 return
-            elif userdata.height == SV(0):
+            if userdata.height == SV(0):
                 return
-            else:
-                userdata.height = smallestsize * Decimal(0.9)
-                userdb.save(userdata)
+            userdata.height = usersizes.smallest.size * Decimal(0.9)
+            userdb.save(userdata)
 
         if lg == m.author.id:
-            if m.author.id == largestuser:
+            if m.author.id == usersizes.largest.id:
                 return
-            elif userdata.height == SV.infinity:
+            if userdata.height == SV.infinity:
                 return
-            else:
-                userdata.height = largestsize * Decimal(1.1)
-                userdb.save(userdata)
+            userdata.height = usersizes.largest.size * Decimal(1.1)
+            userdb.save(userdata)
 
         if userdata.display:
             await nickmanager.nick_update(m.author)