Implement per-user cooldown on Discord

clients/discord.py

@@ -1,4 +1,5 @@
 import re
+import time
 
 import discord
 
@@ -15,6 +16,7 @@ class DiscordClient(discord.Client):
             (re.compile(r'^!(bella(gram|pics)|insta(gram|bella))$'), self._do_bellagram),
             (re.compile(r'^!yt\s+(?P<q>")?(?P<query>.+)(?(q)")$'), self._do_yt),
         ]
+        self.cooldowns = {p.pattern: {} for p, _ in self.supported_commands}
         super(DiscordClient, self).__init__()
 
     async def start_(self):
@@ -25,10 +27,26 @@ class DiscordClient(discord.Client):
         self.logger.info('Logged in as {0}'.format(self.user.name))
 
     async def on_message(self, message):
+        def check_cooldown(pattern, user):
+            retries = self.config['Discord'].getint('cooldown_retries')
+            seconds = self.config['Discord'].getint('cooldown_seconds')
+            if user not in self.cooldowns[pattern]:
+                self.cooldowns[pattern][user] = dict(tries=0, time=0)
+            cooldown = self.cooldowns[pattern][user]
+            if cooldown['tries'] >= retries and time.time() - cooldown['time'] <= seconds:
+                return False
+            cooldown['tries'] = 1 if cooldown['tries'] >= retries else cooldown['tries'] + 1
+            cooldown['time'] = time.time()
+            return True
         for pattern, action in self.supported_commands:
             m = pattern.match(message.content)
             if m:
+                if check_cooldown(pattern.pattern, message.author.id):
                     await action(message, **m.groupdict())
+                else:
+                    await self.send_message(message.channel,
+                        'Sorry {0}, you have to wait a while before running '
+                        'the same command again'.format(message.author.mention))
 
     async def _do_bellagram(self, message, **kwargs):
         try:

settings.cfg.example

@@ -6,6 +6,9 @@ channels = lilialil
 
 [Discord]
 token = __DISCORD_TOKEN__
+# 5 minutes cooldown after 2 retries
+cooldown_retries = 2
+cooldown_seconds = 300
 
 [Twitch]
 api_url = https://api.twitch.tv/v5