Automatically assign team role in 'pick' command

Author: Bluenix2

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

@@ -1,13 +1,18 @@
 from discord import AllowedMentions
 from discord.ext import commands
 from discord.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."""
@@ -21,14 +26,25 @@ async def pick(self, ctx: Context):
         )
 
         index = key % len(constants.ALL_FRUIT_AND_VEG_ROLES)
-        role = constants.ALL_FRUIT_AND_VEG_ROLES[index]
+        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: <@&{role}>",
+            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.add_cog(PickTeam(bot))

tyrant/constants.py

@@ -151,22 +151,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,
 }