Merge pull request #175 from python-discord/levels_init

Add Levels functionality

Author: mbaruh

bot/constants.py

@@ -168,6 +168,17 @@ class _Roles(EnvConfig, env_prefix="ROLE_"):
     team_dict: int = 1222691368653033652
     team_tuple: int = 1222691399246286888
 
+    levels_crystal: int = 1488465184803524758
+    levels_level3: int = 1488465669401088000
+    levels_s_tier: int = 1488465793082462248
+    levels_diamond_rank: int = 1488465915560460399
+    levels_GOAT: int = 1488465529835487232  #noqa: N815
+    levels_mtfn: int = 1488537117976957030
+    levels_champion: int = 1488464806855053433
+    levels_mythical_python_charmer: int = 1488466097412771853
+    levels_supernova_wonder: int = 1488466164106395718
+    levels_ascenion_20: int = 1488466329672351876
+
 
 Roles = _Roles()
 

bot/exts/levels/__init__.py

@@ -0,0 +1,10 @@
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from bot.bot import SirRobin
+
+
+async def setup(bot: "SirRobin") -> None:
+    """Load the CodeJams cog."""
+    from bot.exts.levels._cog import Levels
+    await bot.add_cog(Levels(bot))

bot/exts/levels/_cog.py

@@ -0,0 +1,342 @@
+import random
+import re
+import tomllib
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Literal
+
+import discord
+from async_rediscache import RedisCache
+from discord.ext import commands, tasks
+from pydis_core.utils.logging import get_logger
+
+from bot import constants
+from bot.bot import SirRobin
+from bot.utils import members
+from bot.utils.decorators import in_whitelist
+
+logger = get_logger(__name__)
+
+ELEVATED_ROLES = (constants.Roles.admins, constants.Roles.moderation_team, constants.Roles.events_lead)
+
+ALLOWED_COMMAND_CHANNELS = (constants.Channels.bot_commands, constants.Channels.sir_lancebot_playground,)
+
+# Channels where the game runs.
+ALLOWED_CHANNELS = (
+    constants.Channels.off_topic_0,
+    constants.Channels.off_topic_1,
+    constants.Channels.off_topic_2,
+)
+
+LEVEL_ROLES = (
+    constants.Roles.levels_crystal,
+    constants.Roles.levels_level3,
+    constants.Roles.levels_s_tier,
+    constants.Roles.levels_diamond_rank,
+    constants.Roles.levels_GOAT,
+    constants.Roles.levels_mtfn,
+    constants.Roles.levels_champion,
+    constants.Roles.levels_mythical_python_charmer,
+    constants.Roles.levels_supernova_wonder,
+    constants.Roles.levels_ascenion_20
+)
+
+class Levels(commands.Cog):
+    """Cog that handles all Level functionality."""
+
+    #RedisCache[user_id: int, points: int]
+    user_points_cache = RedisCache()
+
+    #RedisCache[role_id: int, point_threshold: int]
+    levels_cache = RedisCache()
+
+    #RedisCache["value", bool]
+    running = RedisCache()
+
+    def __init__(self, bot: SirRobin):
+        self.bot = bot
+
+        self.rules_all = []
+        self.rules_pool = []
+        self.rules_active = []
+        self.rules_folder_path = Path("./bot/exts/levels/rules/")
+        self.active_rules_num = 3
+
+        self.active_reaction_rule_triggers = []
+        self.active_message_rule_triggers = []
+
+
+    async def cog_load(self) -> None:
+        """Run startup tasks needed when cog is first loaded."""
+        await self._load_rules()
+
+        # Fill in cache with data for later functions to use
+        if await self.levels_cache.length() == 0:
+            shuffled_roles = random.sample(LEVEL_ROLES, len(LEVEL_ROLES))
+            init_threshold_dict = dict.fromkeys(shuffled_roles, 0)
+            await self.levels_cache.update(init_threshold_dict)
+
+        if await self.running.get("value", False):
+            logger.debug("Starting Rules and Point Renormalization tasks")
+            await self._cycle_rules_task.start()
+            await self._calculate_point_thresholds_task.start()
+
+    async def _load_rules(self) -> None:
+        """
+        Load and parse levels rules for usage.
+
+        If a rule file does not comply with the format
+        and throws and error, it is skipped over.
+        """
+        total_files_loaded = 0
+        for toml_file in self.rules_folder_path.glob("*.toml"):
+            with open(toml_file, "rb") as f:
+                rule_dict = tomllib.load(f)
+
+            rule_name = toml_file.stem
+            try:
+                rule_triggers = [RuleTrigger(**rule_trigger) for rule_trigger in rule_dict["rule"]]
+                rule = LevelRules(rule_name, rule_triggers)
+            except (TypeError, KeyError):
+                logger.info(f"{toml_file} not properly formatted, skipping.")
+                continue
+
+            self.rules_all.append(rule)
+            total_files_loaded += 1
+
+        logger.info(f"Total rules loaded: {total_files_loaded}")
+
+    @tasks.loop(minutes=42.0)
+    async def _cycle_rules_task(self) -> None:
+        """
+        Change which rules are currently active.
+
+        Rules will statistically be used before a repeat is seen.
+        This is not a guarnatee though.
+        """
+        if len(self.rules_pool) < self.active_rules_num:
+            # If pool is empty, reshuffle completely to avoid activating same rule twice
+            self.rules_pool = random.sample(self.rules_all, len(self.rules_all))
+        self.rules_active = [self.rules_pool.pop() for _ in range(self.active_rules_num)]
+        logger.debug(f"Cycled active rules to: {[rule.name for rule in self.rules_active]}")
+
+
+        self.active_message_rule_triggers = [
+            rule for rule in self.rules_active
+            for rule_trigger in rule.rule_triggers if rule_trigger.interaction_type=="message"
+        ]
+        self.active_reaction_rule_triggers = [
+            rule for rule in self.rules_active
+            for rule_trigger in rule.rule_triggers if rule_trigger.interaction_type=="reaction"
+        ]
+        # [rule for rule in self.rules_active if rule.interaction_type=="reaction"]
+        # self.active_message_rule_triggers = [rule for rule in self.rules_active if rule.interaction_type=="message"]
+
+    @tasks.loop(minutes=90.0)
+    async def _calculate_point_thresholds_task(self) -> None:
+        """
+        Calculate point thresholds based on number of roles, aiming for even deciles based on scores.
+
+        If current max score is less than 100, it will fix deciles to increments of 10.
+        """
+        user_points = await self.user_points_cache.to_dict()
+        all_scores = sorted(user_points.values())
+        if all_scores and all_scores[-1] >= 100:
+            num_scores = len(all_scores)
+            num_levels = len(LEVEL_ROLES)
+            thresholds = [
+                all_scores[round(num_scores * level/num_levels)]
+                for level in range(1, num_levels+1)
+            ]
+        else:
+            # At the start of the event, just use multiples of 10 up to 100
+            thresholds = [10, 20, 30, 40, 50, 60, 70, 80, 90, 100]
+
+        levels = await self.levels_cache.to_dict()
+        new_levels = dict(zip(levels.keys(), thresholds, strict=False))
+        await self.levels_cache.update(new_levels)
+        logger.debug(f"Renormalizing score thresholds. Total scores: {len(all_scores)}")
+        logger.debug(f"New thresholds: {thresholds}")
+
+
+    async def _update_points(self, user_id: int, points: int) -> None:
+        """Updates user's score and ensures correct role is assigned."""
+        logger.debug(f"User {user_id} getting {points} points.")
+        if not await self.user_points_cache.contains(user_id):
+            await self.user_points_cache.set(user_id, points)
+        else:
+            if points == 0:
+                return
+
+            current_points = await self.user_points_cache.get(user_id)
+            new_point_total = current_points + points
+            await self.user_points_cache.set(user_id, new_point_total)
+
+        await self._update_role_assignment(user_id)
+
+
+    async def _update_role_assignment(self, user_id: int) -> None:
+        """Updates user's role based on current points and role-point thresholds."""
+        user_points = await self.user_points_cache.get(user_id)
+        levels = await self.levels_cache.to_dict()
+        level_to_assign = None
+
+        for role, point_threshold in sorted(levels.items(), key=lambda item: item[1]):
+            level_to_assign = role
+            if point_threshold >= user_points:
+                break
+
+        guild = self.bot.get_guild(constants.Bot.guild)
+        role = guild.get_role(level_to_assign)
+        user = await members.get_or_fetch_member(guild, user_id)
+        if role in user.roles:
+            return
+        logger.debug(f"Assigning {role.name} to {user.name}")
+        await members.handle_role_change(user, user.add_roles, role)
+
+
+    @commands.Cog.listener()
+    async def on_message(self, msg: discord.Message) -> None:
+        """Listens to messages and checks against active message rules."""
+        if not await self.running.get("value", False):
+            return
+        if msg.channel.id not in ALLOWED_CHANNELS or msg.author.bot:
+            return
+        if len(self.active_message_rules) == 0:
+            return
+
+        total_points = 0
+        rule_matches = 0
+        for rule_trigger in self.active_message_rule_triggers:
+            re_pattern = rule_trigger.message_content
+            match = re.search(re_pattern, msg.content)
+            if match:
+                total_points += rule_trigger.points
+                rule_matches += 1
+
+        # Only update points if they've matched any rules
+        # If they match multiple rules and earn 0 points,
+        # that should still get them a role
+        if rule_matches != 0:
+            user_id = msg.author.id
+            await self._update_points(user_id, total_points)
+
+
+    @commands.Cog.listener()
+    async def on_reaction_add(self, reaction: discord.Reaction, user: discord.Member) -> None:
+        """
+        Listens for reactions and checks for against active reaction rules.
+
+        It will only listen for reactions added to messages within the bot's message cache.
+        """
+        if not await self.running.get("value", False):
+            return
+        if reaction.message.channel.id not in ALLOWED_CHANNELS or user.bot:
+            return
+        if len(self.active_reaction_rules) == 0:
+            return
+
+        if isinstance(reaction.emoji, str):
+            emoji_name = reaction.emoji
+        else:
+            emoji_name = reaction.emoji.name
+
+        total_points = 0
+        rule_matches = 0
+        for rule_trigger in self.active_reaction_rule_triggers:
+            if emoji_name in rule_trigger.reaction_content:
+                total_points += rule_trigger.points
+                rule_matches += 1
+
+        # Only update points if they've matched any rules
+        # If they match multiple rules and earn 0 points,
+        # that should still get them a role
+        if rule_matches != 0:
+            await self._update_points(user.id, total_points)
+
+
+    @commands.group(name="levels")
+    async def levels_command_group(self, ctx: commands.Context) -> None:
+        """Levels group command."""
+        if not ctx.invoked_subcommand:
+            await self.bot.invoke_help_command(ctx)
+
+
+    @levels_command_group.command()
+    @in_whitelist(channels=ALLOWED_COMMAND_CHANNELS)
+    async def points(self, ctx: commands.Context) -> None:
+        """Check how many points you've accrued for the Role Level system."""
+        user_id = ctx.author.id
+
+        if await self.user_points_cache.contains(user_id):
+            points = await self.user_points_cache.get(user_id)
+            await ctx.reply(f"You have {points} points.")
+        else:
+            await ctx.reply("You have not earned any points so far! :D")
+
+
+    @levels_command_group.command()
+    @commands.has_any_role(*ELEVATED_ROLES)
+    async def shuffle_role_order(self, ctx: commands.Context) -> None:
+        """Shuffle which roles are assigned to which point thresholds."""
+        levels = await self.levels_cache.to_dict()
+        thresholds = levels.values()
+
+        role_order = random.sample(LEVEL_ROLES, len(LEVEL_ROLES))
+        updated_ordering = dict(zip(role_order, thresholds, strict=False))
+
+        await self.levels_cache.update(updated_ordering)
+        logger.info(f"Roles have been re-shuffled per request of {ctx.author.name}")
+
+    @levels_command_group.command()
+    @commands.has_any_role(*ELEVATED_ROLES)
+    async def start(self, ctx: commands.Context) -> None:
+        """Allows Levels to run, check messages, and assign roles."""
+        current_state = await self.running.get("value", False)
+        if current_state:
+            await ctx.reply("Levels is already running.")
+            return
+
+        self._cycle_rules_task.start()
+        self._calculate_point_thresholds_task.start()
+        await self.running.set("value", True)
+        await ctx.reply("Levels is now turned on.")
+
+    @levels_command_group.command()
+    @commands.has_any_role(*ELEVATED_ROLES)
+    async def stop(self, ctx: commands.Context) -> None:
+        """Disallows Levels to run, check messages, and assign roles."""
+        current_state = await self.running.get("value", False)
+        if not current_state:
+            await ctx.reply("Levels is already off.")
+            return
+
+        self._cycle_rules_task.cancel()
+        self._calculate_point_thresholds_task.cancel()
+        await self.running.set("value", False)
+        await ctx.reply("Levels is now turned off.")
+
+    @levels_command_group.command()
+    @commands.has_any_role(*ELEVATED_ROLES)
+    async def status(self, ctx: commands.Context) -> None:
+        """Replies with current status of Levels."""
+        current_state = await self.running.get("value", False)
+        if current_state:
+            await ctx.reply(":white_check_mark: Levels is currently running.")
+        else:
+            await ctx.reply(":x: Levels is current **not** running.")
+
+
+# Please see ./rules/README.md for how to format rules
+
+@dataclass
+class RuleTrigger:
+    interaction_type: Literal["message", "reaction"]
+    reaction_content: list[str] | None = None
+    message_content: str | None = None
+    points: int = 0
+
+@dataclass
+class LevelRules:
+    name: str
+    rule_triggers: list[RuleTrigger]