Initial move of Snakes cog from the bot to seasonalbot

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