Skip to content

Commit 49d60ac

Browse files
committed
allow user specified function to manipulate cooldown args
1 parent 2c903cf commit 49d60ac

File tree

1 file changed

+9
-3
lines changed

1 file changed

+9
-3
lines changed

botcore/utils/cooldown.py

Lines changed: 9 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -175,13 +175,16 @@ def block_duplicate_invocations(
175175
*,
176176
cooldown_duration: float = 5,
177177
send_notice: bool = False,
178+
args_preprocessor: Callable[P, Iterable[object]] | None = None,
178179
) -> Callable[[Callable[P, Awaitable[R]]], Callable[P, Awaitable[R]]]:
179180
"""
180181
Prevent duplicate invocations of a command with the same arguments in a channel for ``cooldown_duration`` seconds.
181182
182183
Args:
183184
cooldown_duration: Length of the cooldown in seconds.
184185
send_notice: If :obj:`True`, notify the user about the cooldown with a reply.
186+
args_preprocessor: If specified, this function is called with the args and kwargs the function is called with,
187+
its return value is then used to check for the cooldown instead of the raw arguments.
185188
186189
Returns:
187190
A decorator that adds a wrapper which applies the cooldowns.
@@ -195,16 +198,19 @@ def decorator(func: Callable[P, Awaitable[R]]) -> Callable[P, Awaitable[R]]:
195198

196199
@command_wraps(func)
197200
async def wrapper(*args: P.args, **kwargs: P.kwargs) -> R:
198-
arg_tuple = _create_argument_tuple(*args[2:], **kwargs) # skip self and ctx from the command
201+
if args_preprocessor is not None:
202+
all_args = args_preprocessor(*args, **kwargs)
203+
else:
204+
all_args = _create_argument_tuple(*args[2:], **kwargs) # skip self and ctx from the command
199205
ctx = typing.cast("Context[BotBase]", args[1])
200206

201207
if not isinstance(ctx.channel, discord.DMChannel):
202-
if mgr.is_on_cooldown(ctx.channel, arg_tuple):
208+
if mgr.is_on_cooldown(ctx.channel, all_args):
203209
if send_notice:
204210
with suppress(discord.NotFound):
205211
await ctx.reply("The command is on cooldown with the given arguments.")
206212
raise CommandOnCooldown(ctx.message.content, func, *args, **kwargs)
207-
mgr.set_cooldown(ctx.channel, arg_tuple)
213+
mgr.set_cooldown(ctx.channel, all_args)
208214

209215
return await func(*args, **kwargs)
210216

0 commit comments

Comments
 (0)