Make cooldowns per-server

clients/discord.py

@@ -8,22 +8,23 @@ from commands import CommandError
 
 def cooldown(retries, timeout, failure):
     def do_cooldown(function):
-        def wrapper(self, user, *args, **kwargs):
+        def wrapper(self, server, user, *args, **kwargs):
             cooldowns = getattr(function, 'cooldowns', {})
-            if user not in cooldowns:
+            hash = '{0}:{1}'.format(server, user)
-                cooldowns[user] = dict(tries=0, last_run=0)
+            if hash not in cooldowns:
-            cd = cooldowns[user]
+                cooldowns[hash] = dict(tries=0, last_run=0)
+            cd = cooldowns[hash]
             if cd['tries'] < retries:
                 cd['tries'] += 1
                 cd['last_run'] = time.time()
                 setattr(function, 'cooldowns', cooldowns)
-                return function(self, user, *args, **kwargs)
+                return function(self, server, user, *args, **kwargs)
             if time.time() - cd['last_run'] > timeout:
                 cd['tries'] = 1
                 cd['last_run'] = time.time()
                 setattr(function, 'cooldowns', cooldowns)
-                return function(self, user, *args, **kwargs)
+                return function(self, server, user, *args, **kwargs)
-            return failure(self, user, *args, **kwargs)
+            return failure(self, server, user, *args, **kwargs)
         return wrapper
     return do_cooldown
 
@@ -52,15 +53,16 @@ class DiscordClient(discord.Client):
         for pattern, action in self.supported_commands:
             m = pattern.match(message.content)
             if m:
-                await action(message.author.id, message, **m.groupdict())
+                server = message.server.id if message.server else None
+                await action(server, message.author.id, message, **m.groupdict())
 
-    async def _cooldown_failure(self, user, message, **kwargs):
+    async def _cooldown_failure(self, server, user, message, **kwargs):
         await self.send_message(message.channel,
             'Sorry {0}, you have to wait a while before running '
             'the same command again'.format(message.author.mention))
 
     @cooldown(retries=2, timeout=5*60, failure=_cooldown_failure)
-    async def _do_bellagram(self, user, message, **kwargs):
+    async def _do_bellagram(self, server, user, message, **kwargs):
         try:
             bellagram = self.commands.bellagram()
         except CommandError as e:
@@ -73,7 +75,7 @@ class DiscordClient(discord.Client):
             await self.send_message(message.channel, embed=embed)
 
     @cooldown(retries=3, timeout=5*60, failure=_cooldown_failure)
-    async def _do_yt(self, user, message, query, **kwargs):
+    async def _do_yt(self, server, user, message, query, **kwargs):
         try:
             result = self.commands.query_youtube(query)
         except CommandError as e:
@@ -90,7 +92,7 @@ class DiscordClient(discord.Client):
                 await self.send_message(message.channel, '{title}\n{url}'.format(**result))
 
     @cooldown(retries=3, timeout=5*60, failure=_cooldown_failure)
-    async def _do_clip(self, user, message, filter, **kwargs):
+    async def _do_clip(self, server, user, message, filter, **kwargs):
         try:
             result = self.commands.find_clip(filter)
         except CommandError as e: