diff options
Diffstat (limited to 'bot/decorators.py')
-rw-r--r-- | bot/decorators.py | 41 |
1 files changed, 41 insertions, 0 deletions
diff --git a/bot/decorators.py b/bot/decorators.py index b84b2c36..f5ffadf4 100644 --- a/bot/decorators.py +++ b/bot/decorators.py @@ -1,8 +1,15 @@ import logging +import random +from asyncio import Lock +from functools import wraps +from weakref import WeakValueDictionary +from discord import Colour, Embed from discord.ext import commands from discord.ext.commands import Context +from bot.constants import ERROR_REPLIES + log = logging.getLogger(__name__) @@ -46,3 +53,37 @@ def in_channel(channel_id): f"The result of the in_channel check was {check}.") return check return commands.check(predicate) + + +def locked(): + """ + Allows the user to only run one instance of the decorated command at a time. + Subsequent calls to the command from the same author are + ignored until the command has completed invocation. + + This decorator has to go before (below) the `command` decorator. + """ + + def wrap(func): + func.__locks = WeakValueDictionary() + + @wraps(func) + async def inner(self, ctx, *args, **kwargs): + lock = func.__locks.setdefault(ctx.author.id, Lock()) + if lock.locked(): + embed = Embed() + embed.colour = Colour.red() + + log.debug(f"User tried to invoke a locked command.") + embed.description = ( + "You're already using this command. Please wait until " + "it is done before you use it again." + ) + embed.title = random.choice(ERROR_REPLIES) + await ctx.send(embed=embed) + return + + async with func.__locks.setdefault(ctx.author.id, Lock()): + return await func(self, ctx, *args, **kwargs) + return inner + return wrap |