Add command that let's Tyrant pick a user's role

tyrant/__main__.py

@@ -19,10 +19,11 @@
     intents=intents,
 )
 
-# Load the extensions we want
+# Load the extensions we want (alphabetical)
 bot.load_extension("tyrant.cogs.ask_tyrant")
-bot.load_extension("tyrant.cogs.lemon_facts")
 bot.load_extension("tyrant.cogs.fruit_vs_vegetables")
+bot.load_extension("tyrant.cogs.lemon_facts")
+bot.load_extension("tyrant.cogs.pick_team")
 bot.load_extension("tyrant.cogs.purge")
 bot.load_extension("tyrant.cogs.teamcount")
 

tyrant/cogs/fruit_vs_vegetables.py

@@ -13,6 +13,11 @@ def __init__(self, bot: Bot):
         self.bot = bot
         self.locks = {}
 
+    async def assign_roles(self, member, *roles) -> None:
+        """Assign roles to a member and remove the old team roles."""
+        await member.remove_roles(*[role for role in member.roles if role.id in constants.ALL_FRUIT_AND_VEG_ROLES])
+        await member.add_roles(*roles)
+
     @Cog.listener()
     async def on_raw_reaction_add(self, payload):
         """Distribute fruit or vegetable role, when appropriate."""
@@ -36,13 +41,11 @@ async def on_raw_reaction_add(self, payload):
 
                 # Get the role ID from the emoji
                 fruit_role_id = constants.EMOJI_TO_ROLE[emoji.name]
-                team_id = constants.EMOJI_TO_TEAM[emoji.name]
+                team_id = constants.ROLE_TO_TEAM[fruit_role_id]
                 fruit_role = guild.get_role(fruit_role_id)
                 team_role = guild.get_role(team_id)
 
-                # Get rid of old roles, assign the new ones
-                await member.remove_roles(*[role for role in member.roles if role.id in constants.ALL_FRUIT_AND_VEG_ROLES])
-                await member.add_roles(fruit_role, team_role)
+                await self.assign_roles(member, fruit_role, team_role)
 
                 # Finally, remove all other reactions than this one
                 fruit_message = await channel.fetch_message(constants.Messages.fruit_role_assignment)
@@ -79,7 +82,7 @@ async def on_raw_reaction_remove(self, payload):
 
                 # Get the role ID from the emoji
                 fruit_role_id = constants.EMOJI_TO_ROLE[emoji.name]
-                team_id = constants.EMOJI_TO_TEAM[emoji.name]
+                team_id = constants.ROLE_TO_TEAM[fruit_role_id]
                 team_role = guild.get_role(team_id)
 
                 # Remove all fruit and veg roles from the member

tyrant/cogs/pick_team.py

@@ -0,0 +1,53 @@
+from disnake import AllowedMentions
+from disnake.ext import commands
+from disnake.ext.commands import Bot, Cog, Context
+from loguru import logger
+
+from tyrant import constants
+
+
+class PickTeam(Cog):
+    """Command letting the Tyrant pick a team for the user."""
+
+    def __init__(self, bot: Bot):
+        """Initialize this cog with the Bot instance."""
+        self.bot = bot
+
+    @commands.command(aliases=("team",))
+    async def pick(self, ctx: Context):
+        """Let the Tyrant pick a fruit or vegetable for you."""
+        if ctx.guild is None:
+            return await ctx.send('This command can only be used in the server.')
+
+        # We want a sense of randomness from this algorithm so that the chosen
+        # roles are evenly spread out, at the same time we want the command to
+        # choose the same roles even if invoked twice.
+        key = hash(
+            ctx.author.display_name +
+            ctx.author.discriminator +
+            ctx.author.display_avatar.key
+        )
+
+        index = key % len(constants.ALL_FRUIT_AND_VEG_ROLES)
+        choice = constants.ALL_FRUIT_AND_VEG_ROLES[index]
+
+        if cog := self.bot.get_cog("FruitVsVegetables"):
+            await cog.assign_roles(
+                ctx.author,
+                ctx.guild.get_role(choice),
+                ctx.guild.get_role(constants.ROLE_TO_TEAM[choice])
+            )
+        else:
+            # The FruitVsVegetables cog being loaded is a pretty important
+            # aspect of this command but we can still recover without it
+            logger.warning('Could not assign role because FruitVsVegetables cog is unloaded')
+
+        await ctx.send(
+            f"**The Tyrant decided that your role will be**... <@&{choice}>",
+            allowed_mentions=AllowedMentions.none()
+        )
+
+
+def setup(bot: Bot) -> None:
+    """Called by discord.py to load the extension which will add the above cog."""
+    bot.add_cog(PickTeam(bot))

tyrant/constants.py

@@ -173,22 +173,22 @@
     "🥑": 883078311494963210,
 }
 
-EMOJI_TO_TEAM = {
-    "🌶️": 883078871770071051,
-    "🥕": 883078871770071051,
-    "yellow_bell_pepper": 883078871770071051,
-    "🥑": 883078871770071051,
-    "blue_asparagus": 883078871770071051,
-    "cabbage": 883078871770071051,
-    "radish": 883078871770071051,
-
-    "pomegranate": 883080439231840259,
-    "🍊": 883080439231840259,
-    "🍌": 883080439231840259,
-    "🥝": 883080439231840259,
-    "🫐": 883080439231840259,
-    "🍇": 883080439231840259,
-    "dragonfruit": 883080439231840259,
+ROLE_TO_TEAM = {
+    883077551927484476: 883078871770071051,
+    883077759985930290: 883078871770071051,
+    883078108457083010: 883078871770071051,
+    883078311494963210: 883078871770071051,
+    883078590655258674: 883078871770071051,
+    882529427962069047: 883078871770071051,
+    802889394074681426: 883078871770071051,
+
+    802850922375282698: 883080439231840259,
+    802843500638240789: 883080439231840259,
+    843204335793602601: 883080439231840259,
+    802851237837799455: 883080439231840259,
+    882360895462846545: 883080439231840259,
+    882633463432355882: 883080439231840259,
+    816010822777831444: 883080439231840259,
 }