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 | 
