diff --git a/tyrant/__main__.py b/tyrant/__main__.py index 9e5ff77..ad0d585 100644 --- a/tyrant/__main__.py +++ b/tyrant/__main__.py @@ -19,11 +19,12 @@ 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.purge") +bot.load_extension("tyrant.cogs.random_team") bot.load_extension("tyrant.cogs.teamcount") # Validate the token diff --git a/tyrant/cogs/fruit_vs_vegetables.py b/tyrant/cogs/fruit_vs_vegetables.py old mode 100755 new mode 100644 index f73226d..7f5f5a2 --- a/tyrant/cogs/fruit_vs_vegetables.py +++ b/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 diff --git a/tyrant/cogs/random_team.py b/tyrant/cogs/random_team.py new file mode 100644 index 0000000..42e49bd --- /dev/null +++ b/tyrant/cogs/random_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 RandomTeam(Cog): + """Command letting the Tyrant pick a random team for the user.""" + + def __init__(self, bot: Bot): + """Initialize this cog with the Bot instance.""" + self.bot = bot + + @commands.command(aliases=("enlist", "enroll")) + async def random_team(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(RandomTeam(bot)) diff --git a/tyrant/constants.py b/tyrant/constants.py index d0c264e..42f8cb1 100644 --- a/tyrant/constants.py +++ b/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, }