aboutsummaryrefslogtreecommitdiffstats
path: root/bot/decorators.py
diff options
context:
space:
mode:
Diffstat (limited to 'bot/decorators.py')
-rw-r--r--bot/decorators.py41
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