diff --git a/tux/cogs/services/gif_limiter.py b/tux/cogs/services/gif_limiter.py index e87f580..21d6bfd 100644 --- a/tux/cogs/services/gif_limiter.py +++ b/tux/cogs/services/gif_limiter.py @@ -1,12 +1,16 @@ +""" +This cog is a handler for GIF ratelimiting. +It keeps a list of GIF send times and routinely removes old times. +If a user posts a GIF, the message_handler function should be externally called. +It will delete the message if the user or channel quota is exceeded. +""" + import asyncio from collections import defaultdict from time import time -from typing import defaultdict - -from tux.utils.constants import Constants as CONST -from tux.bot import Tux import discord +from discord import Message from discord.ext import commands, tasks from loguru import logger @@ -14,8 +18,9 @@ from tux.bot import Tux from tux.utils.constants import CONST -# Helper function required as YAML keys are str. Channel and user IDs are int. def convert_dict_str_to_int(original_dict: dict[str, int]) -> dict[int, int]: + """Helper function required as YAML keys are str. Channel and user IDs are int.""" + converted_dict: dict[int, int] = {} for key, value in original_dict.items(): @@ -29,44 +34,44 @@ def convert_dict_str_to_int(original_dict: dict[str, int]) -> dict[int, int]: class GifLimiter(commands.Cog): - """ - This class is a handler for GIF ratelimiting. - It keeps a list of GIF send times and routinely removes old times. - If a user posts a GIF, the message_handler function should be externally called. - It will delete the message if the user, channel or server-wide quota is exceeded. - """ + """Main class with GIF tracking and message handlers""" def __init__(self, bot: Tux) -> None: self.bot = bot - # Read config options and save them to local variables to avoid excessive reads - # From the CONST[] dictionary - self.recent_gif_age: int = CONST.RECENT_GIF_AGE # Max age for a GIF to be considered a recent post - self.channelwide_gif_limits: dict[int, int] = convert_dict_str_to_int( - CONST.GIF_LIMITS_CHANNEL - ) # Max GIFs sent recently for specific channels - self.user_gif_limits: dict[int, int] = convert_dict_str_to_int( - CONST.GIF_LIMITS - ) # Max recent GIFs sent by a user to be able to send a GIF in a channel - self.gif_limit_exclude: list[int] = CONST.GIF_LIMIT_EXCLUDE # list of channels in which not to count GIFs + + # Max age for a GIF to be considered a recent post + self.recent_gif_age: int = CONST.RECENT_GIF_AGE + + # Max number of GIFs sent recently in a channel + self.channelwide_gif_limits: dict[int, int] = convert_dict_str_to_int(CONST.GIF_LIMITS_CHANNEL) + # Max number of GIFs sent recently by a user to be able to post one in specified channels + self.user_gif_limits: dict[int, int] = convert_dict_str_to_int(CONST.GIF_LIMITS) + + # list of channels in which not to count GIFs + self.gif_limit_exclude: list[int] = CONST.GIF_LIMIT_EXCLUDE # Timestamps for recently-sent GIFs for the server, and channels - self.recent_gifs_by_user: defaultdict[int, list[int]] = defaultdict(list) # UID, list of timestamps - self.recent_gifs_by_channel: defaultdict[int, list[int]] = defaultdict(list) # Channel ID, list of timestamps - self.recent_gifs_serverwide: list[int] = [] - # Deletes the message passed as an argument, and creates a self-deleting message explaining the reason + # UID, list of timestamps + self.recent_gifs_by_user: defaultdict[int, list[int]] = defaultdict(list) + # Channel ID, list of timestamps + self.recent_gifs_by_channel: defaultdict[int, list[int]] = defaultdict(list) + async def delete_message(self, message: discord.Message, epilogue: str) -> None: - channel: Union[TextChannel, StageChannel, VoiceChannel, Thread, DMChannel, GroupChannel, PartialMessageable] - = message.channel - + """ + Deletes the message passed as an argument, and sends a self-deleting message with the reason + """ + sent_message: Message = await message.channel.send(f"-# GIF ratelimit exceeded {epilogue}") await message.delete() - sent_message: discord.Message = await channel.send("-# GIF ratelimit exceeded " + epilogue) await asyncio.sleep(3) await sent_message.delete() @commands.Cog.listener() async def on_message(self, message: discord.Message) -> None: - # Nothing to do if the message doesn't have a .gif embed, or if it was sent in a blacklisted channel + """Checks for GIFs in every sent message""" + + # Nothing to do if the message doesn't have a .gif embed, + # or if it was sent in a blacklisted channel if ( len(message.embeds) == 0 or "gif" not in message.content.lower() @@ -99,16 +104,18 @@ class GifLimiter(commands.Cog): self.recent_gifs_by_channel[channel].append(current_time) self.recent_gifs_by_user[user].append(current_time) - # Function regularly cleans GIF lists and only keeps the most recent ones @tasks.loop(seconds=20) async def old_gif_remover(self) -> None: + """Regularly cleans old GIF timestamps""" current_time: int = int(time()) for channel_id, timestamps in self.recent_gifs_by_channel.items(): - self.recent_gifs_by_channel[channel_id] = [t for t in timestamps if current_time - t < self.recent_gif_age] + self.recent_gifs_by_channel[channel_id] = ( + [t for t in timestamps if current_time - t < self.recent_gif_age]) for user_id, timestamps in self.recent_gifs_by_user.items(): - self.recent_gifs_by_user[user_id] = [t for t in timestamps if current_time - t < self.recent_gif_age] + self.recent_gifs_by_user[user_id] = ( + [t for t in timestamps if current_time - t < self.recent_gif_age]) # Delete user key if no GIF has recently been sent by them if len(self.recent_gifs_by_user[user_id]) == 0: