Add /heatmap command

README.md

@@ -37,7 +37,7 @@ restrict_channel = "new-user"
 welcome_channel = "off-topic"
 
 [Blossom]
-username = ""
+email = ""
 password = ""
 api_key = ""
 ```

buttercup/bot.py

@@ -24,7 +24,7 @@ def __init__(self, command_prefix: str, **kwargs: Any) -> None:
         intents.members = True
         super().__init__(command_prefix, intents=intents, **kwargs)
         self.slash = SlashCommand(self, sync_commands=True, sync_on_cog_reload=True)
-        self.config_path = kwargs.get("config_path", "config.toml")
+        self.config_path = kwargs.get("config_path", "../config.toml")
         self.cog_path = kwargs.get("cog_path", "buttercup.cogs.")
         self.guild_name = self.config["guild"]["name"]
 

buttercup/cogs/__init__.py

@@ -1 +1,20 @@
 """The cogs which provide the functionality to the bot."""
+
+# Colors to use in the plots
+from matplotlib import pyplot as plt
+
+background_color = "#36393f"  # Discord background color
+text_color = "white"
+line_color = "white"
+
+# Global settings for all plots
+plt.rcParams["figure.facecolor"] = background_color
+plt.rcParams["axes.facecolor"] = background_color
+plt.rcParams["axes.labelcolor"] = text_color
+plt.rcParams["axes.edgecolor"] = line_color
+plt.rcParams["text.color"] = text_color
+plt.rcParams["xtick.color"] = line_color
+plt.rcParams["ytick.color"] = line_color
+plt.rcParams["grid.color"] = line_color
+plt.rcParams["grid.alpha"] = 0.8
+plt.rcParams["figure.dpi"] = 200.0

buttercup/cogs/handlers.py

@@ -5,6 +5,10 @@
 
 from buttercup import logger
 from buttercup.bot import ButtercupBot
+from buttercup.cogs.helpers import NoUsernameError
+from buttercup.strings import translation
+
+i18n = translation()
 
 
 class Handlers(commands.Cog):
@@ -25,17 +29,14 @@ async def on_slash_command_error(
         """Log that a command has errored and provide the user with feedback."""
         if isinstance(error, commands.CheckFailure):
             logger.warning("An unauthorized Command was performed.", ctx)
-            await ctx.send(
-                "You are not authorized to use this command. "
-                "This incident will be reported"
-            )
+            await ctx.send(i18n["handlers"]["not_authorized"])
+        elif isinstance(error, NoUsernameError):
+            logger.warning("Command executed without providing a username.", ctx)
+            await ctx.send(i18n["handlers"]["no_username"])
         else:
             tracker_id = uuid.uuid4()
             logger.warning(f"[{tracker_id}] {type(error).__name__}: {str(error)}", ctx)
-            await ctx.send(
-                f"[{tracker_id}] Something went wrong, "
-                "please contact a moderator with the provided ID."
-            )
+            await ctx.send(i18n["handlers"]["unknown_error"].format(tracker_id))
 
 
 def setup(bot: ButtercupBot) -> None:

buttercup/cogs/heatmap.py

@@ -0,0 +1,147 @@
+import io
+from datetime import datetime
+from typing import Any, Dict, Optional
+
+import matplotlib.pyplot as plt
+import pandas as pd
+import seaborn as sns
+from blossom_wrapper import BlossomAPI
+from discord import File
+from discord.ext.commands import Cog
+from discord_slash import SlashContext, cog_ext
+from discord_slash.utils.manage_commands import create_option
+
+from buttercup.bot import ButtercupBot
+from buttercup.cogs.helpers import (
+    extract_username,
+    extract_utc_offset,
+    get_duration_str,
+)
+from buttercup.strings import translation
+
+i18n = translation()
+
+
+def create_file_from_heatmap(
+    heatmap: pd.DataFrame, username: str, utc_offset: int = 0
+) -> File:
+    """Create a Discord file containing the heatmap table."""
+    days = i18n["heatmap"]["days"]
+    hours = ["{:02d}".format(hour) for hour in range(0, 24)]
+
+    # The built in formatting for the heatmap doesn't allow displaying floats as ints
+    # And we have to use floats because empty entries are NaN
+    # So we have to manually provide the annotations
+    annotations = heatmap.apply(
+        lambda series: series.apply(lambda value: f"{value:0.0f}")
+    )
+
+    fig, ax = plt.subplots()
+    fig.set_size_inches(9, 3.44)
+
+    sns.heatmap(
+        heatmap,
+        ax=ax,
+        annot=annotations,
+        fmt="s",
+        cbar=False,
+        square=True,
+        xticklabels=hours,
+        yticklabels=days,
+    )
+
+    timezone = "UTC" if utc_offset == 0 else f"UTC{utc_offset:+d}"
+
+    plt.title(i18n["heatmap"]["plot_title"].format(username))
+    plt.xlabel(i18n["heatmap"]["plot_xlabel"].format(timezone))
+    plt.ylabel(i18n["heatmap"]["plot_ylabel"])
+
+    fig.tight_layout()
+    heatmap_table = io.BytesIO()
+    plt.savefig(heatmap_table, format="png")
+    heatmap_table.seek(0)
+    plt.clf()
+
+    return File(heatmap_table, "heatmap_table.png")
+
+
+def adjust_with_timezone(hour_data: Dict[str, Any], utc_offset: int) -> Dict[str, Any]:
+    """Adjust the heatmap data according to the UTC offset of the user."""
+    hour_offset = hour_data["hour"] + utc_offset
+    new_hour = hour_offset % 24
+    # The days go from 1 to 7, so we need to adjust this to zero index and back
+    new_day = ((hour_data["day"] + hour_offset // 24) - 1) % 7 + 1
+    return {"day": new_day, "hour": new_hour, "count": hour_data["count"]}
+
+
+class Heatmap(Cog):
+    def __init__(self, bot: ButtercupBot, blossom_api: BlossomAPI) -> None:
+        """Initialize the Heatmap cog."""
+        self.bot = bot
+        self.blossom_api = blossom_api
+
+    @cog_ext.cog_slash(
+        name="heatmap",
+        description="Display the activity heatmap for the given user.",
+        options=[
+            create_option(
+                name="username",
+                description="The user to get the heatmap for.",
+                option_type=3,
+                required=False,
+            )
+        ],
+    )
+    async def _heatmap(self, ctx: SlashContext, username: Optional[str] = None) -> None:
+        """Generate a heatmap for the given user."""
+        start = datetime.now()
+        user = username or extract_username(ctx.author.display_name)
+        utc_offset = extract_utc_offset(ctx.author.display_name)
+        msg = await ctx.send(i18n["heatmap"]["getting_heatmap"].format(user))
+
+        response = self.blossom_api.get("volunteer/heatmap/", params={"username": user})
+
+        if response.status_code != 200:
+            await msg.edit(content=i18n["heatmap"]["user_not_found"].format(user))
+            return
+
+        data = response.json()
+        data = [adjust_with_timezone(hour_data, utc_offset) for hour_data in data]
+
+        day_index = pd.Index(range(1, 8))
+        hour_index = pd.Index(range(0, 24))
+
+        heatmap = (
+            # Create a data frame from the data
+            pd.DataFrame.from_records(data)
+            # Convert it into a table with the days as rows and hours as columns
+            .pivot(index="day", columns="hour", values="count")
+            # Add the missing days and hours
+            .reindex(index=day_index, columns=hour_index)
+        )
+
+        heatmap_table = create_file_from_heatmap(heatmap, user, utc_offset)
+
+        await msg.edit(
+            content=i18n["heatmap"]["response_message"].format(
+                user, get_duration_str(start)
+            ),
+            file=heatmap_table,
+        )
+
+
+def setup(bot: ButtercupBot) -> None:
+    """Set up the Heatmap cog."""
+    # Initialize blossom api
+    cog_config = bot.config["Blossom"]
+    email = cog_config.get("email")
+    password = cog_config.get("password")
+    api_key = cog_config.get("api_key")
+    blossom_api = BlossomAPI(email=email, password=password, api_key=api_key)
+
+    bot.add_cog(Heatmap(bot=bot, blossom_api=blossom_api))
+
+
+def teardown(bot: ButtercupBot) -> None:
+    """Unload the Heatmap cog."""
+    bot.remove_cog("Heatmap")

buttercup/cogs/helpers.py

@@ -0,0 +1,36 @@
+import re
+from datetime import datetime
+
+from discord import DiscordException
+
+username_regex = re.compile(r"^/?u/(?P<username>\S+)")
+timezone_regex = re.compile(r"UTC(?P<offset>[+-]\d+)?", re.RegexFlag.I)
+
+
+class NoUsernameError(DiscordException):
+    """Exception raised when the username was not provided."""
+
+    pass
+
+
+def extract_username(display_name: str) -> str:
+    """Extract the Reddit username from the display name."""
+    match = username_regex.search(display_name)
+    if match is None:
+        raise NoUsernameError()
+    return match.group("username")
+
+
+def extract_utc_offset(display_name: str) -> int:
+    """Extract the user's timezone (UTC offset) from the display name."""
+    match = timezone_regex.search(display_name)
+    if match is None or match.group("offset") is None:
+        return 0
+    return int(match.group("offset"))
+
+
+def get_duration_str(start: datetime) -> str:
+    """Get the processing duration based on the start time."""
+    duration = datetime.now() - start
+    duration_ms = duration.microseconds / 1000
+    return f"{duration_ms:0.0f} ms"

buttercup/main.py

@@ -3,7 +3,16 @@
 from buttercup import logger
 from buttercup.bot import ButtercupBot
 
-EXTENSIONS = ["admin", "handlers", "name_validator", "find", "stats", "rules", "ping"]
+EXTENSIONS = [
+    "admin",
+    "handlers",
+    "name_validator",
+    "find",
+    "stats",
+    "heatmap",
+    "ping",
+    "rules",
+]
 
 
 logger.configure_logging()

buttercup/strings/en_US.yaml

@@ -1,3 +1,10 @@
+handlers:
+  not_authorized: |
+    You are not authorized to use this command. This incident will be reported.
+  no_username: |
+    You did not provide a valid username. Either add it to the command or change your display name to a valid format.
+  unknown_error: |
+    [{0}] Something went wrong, please contact a moderator with the provided ID.
 name_validator:
   new_member: |
     Hello, <@{0}>! Welcome to the Transcribers of Reddit chill room! Thanks for helping make somebody's day better!
@@ -36,6 +43,27 @@ stats:
     **Volunteers**: {:,d}
     **Transcriptions**: {:,d}
     **Days Since Inception**: {:,d}
+heatmap:
+  getting_heatmap: |
+    Generating a heatmap for u/{0}...
+  user_not_found: |
+    Sorry, I couldn't find the user u/{0}.
+  days:
+    - Mon
+    - Tue
+    - Wed
+    - Thu
+    - Fri
+    - Sat
+    - Sun
+  plot_title: |
+    Activity Heatmap for u/{0}
+  plot_xlabel: |
+    Time ({0})
+  plot_ylabel: |
+    Weekday
+  response_message: |
+    Here is the heatmap for u/{0}! ({1})
 rules:
   getting_rules: |
     Getting the rules for r/{0}...