From 8c0d98a4a8cc819ddbe1868541b06f027660f4e8 Mon Sep 17 00:00:00 2001 From: Natalie Fearnley Date: Tue, 21 May 2024 02:01:56 -0400 Subject: [PATCH] cleaned up edge.py --- sizebot/cogs/edge.py | 100 +++++++++++++++++++++++-------------------- 1 file changed, 53 insertions(+), 47 deletions(-) diff --git a/sizebot/cogs/edge.py b/sizebot/cogs/edge.py index fbde5f94..52a9d9b9 100644 --- a/sizebot/cogs/edge.py +++ b/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)