shaiwatcher/modules/usage/usage_stats.py
Franz Rolfsvaag fe09e1dd1f 0.4.1.0.a5
- Yet another patch for the exectuion counter..
  - *When stuff won't play nice, you brute-force it*
2025-08-16 02:40:38 +02:00

67 lines
2.7 KiB
Python

# modules/usage/usage_stats.py
from __future__ import annotations
from discord.ext import commands
import discord
COUNTER_KEY_PREFIX = "cmd::"
LOCK_WINDOW_SEC = 1.0 # timelock window; change if you want stricter/looser locking
def _key_from_app(cmd: discord.app_commands.Command) -> str:
name = getattr(cmd, "qualified_name", None) or getattr(cmd, "name", "unknown")
return f"{COUNTER_KEY_PREFIX}{name}"
def _key_from_ctx(ctx: commands.Context) -> str:
c = getattr(ctx, "command", None)
name = getattr(c, "qualified_name", None) or getattr(c, "name", "unknown")
return f"{COUNTER_KEY_PREFIX}{name}"
class UsageStatsCog(commands.Cog):
"""Command run counters with a per-command timelock."""
def __init__(self, bot: commands.Bot):
self.bot = bot
print("[usage] UsageStatsCog init (timelock)")
@commands.Cog.listener()
async def on_app_command_completion(self, interaction: discord.Interaction, command: discord.app_commands.Command):
dm = getattr(self.bot, "data_manager", None)
if not dm:
return
try:
counter_key = _key_from_app(command)
newv = dm.incr_counter_timelocked(counter_key, window_sec=LOCK_WINDOW_SEC)
if newv is not None:
print(f"[usage] app ++ {counter_key} -> {newv}")
else:
print(f"[usage] app ~~ timelocked {counter_key}")
except Exception as e:
print("[usage] app !! incr failed:", repr(e))
@commands.Cog.listener()
async def on_command_completion(self, ctx: commands.Context):
# If a HybridCommand was invoked as a slash interaction, let the app listener count it.
if isinstance(getattr(ctx, "command", None), commands.HybridCommand) and getattr(ctx, "interaction", None):
return
dm = getattr(self.bot, "data_manager", None)
if not dm:
return
try:
counter_key = _key_from_ctx(ctx)
newv = dm.incr_counter_timelocked(counter_key, window_sec=LOCK_WINDOW_SEC)
if newv is not None:
print(f"[usage] px ++ {counter_key} -> {newv}")
else:
print(f"[usage] px ~~ timelocked {counter_key}")
except Exception as e:
print("[usage] px !! incr failed:", repr(e))
async def setup(bot: commands.Bot):
# Prevent duplicate registration if extensions are reloaded / auto-discovered twice
if getattr(bot, "_usage_stats_loaded", False):
print("[usage] UsageStatsCog already loaded; skipping duplicate add")
return
await bot.add_cog(UsageStatsCog(bot))
bot._usage_stats_loaded = True
print("[usage] UsageStatsCog loaded (timelock)")