bot/clients/discord.py

import re
import time

import discord

from commands import CommandError


def cooldown(retries, timeout, failure):
    def do_cooldown(function):
        def wrapper(self, server, user, *args, **kwargs):
            cooldowns = getattr(function, 'cooldowns', {})
            hash = '{0}:{1}'.format(server, user)
            if hash not in cooldowns:
                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, 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, server, user, *args, **kwargs)
            return failure(self, server, user, *args, **kwargs)
        return wrapper
    return do_cooldown


class DiscordClient(discord.Client):

    def __init__(self, config, logger, commands):
        self.config = config
        self.logger = logger
        self.commands = commands
        self.supported_commands = [
            (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),
            (re.compile(r'^!clip\s+(?P<q>")?(?P<filter>.+)(?(q)")$'), self._do_clip),
        ]
        super(DiscordClient, self).__init__()

    async def start_(self):
        token = self.config['Discord'].get('token')
        await self.start(token)

    async def on_ready(self):
        self.logger.info('Logged in as {0}'.format(self.user.name))

    async def on_message(self, message):
        for pattern, action in self.supported_commands:
            m = pattern.match(message.content)
            if m:
                server = message.server.id if message.server else None
                await action(server, message.author.id, message, **m.groupdict())

    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, server, user, message, **kwargs):
        try:
            bellagram = self.commands.bellagram()
        except CommandError as e:
            await self.send_message(message.channel, 'Sorry {0}, {1}'.format(message.author.mention, e))
        else:
            embed = discord.Embed(title=bellagram['title'], url=bellagram['url'], color=0xd2a517)
            embed.set_image(url=bellagram['display_url'])
            embed.set_author(name=bellagram['owner'], url=bellagram['owner_url'],
                icon_url=bellagram['owner_pic_url'])
            await self.send_message(message.channel, embed=embed)

    @cooldown(retries=3, timeout=5*60, failure=_cooldown_failure)
    async def _do_yt(self, server, user, message, query, **kwargs):
        try:
            result = self.commands.query_youtube(query)
        except CommandError as e:
            await self.send_message(message.channel, 'Sorry {0}, {1}'.format(message.author.mention, e))
        else:
            if result['kind'] == 'playlist':
                embed = discord.Embed(title=result['title'], url=result['url'],
                    description=result['description'], color=0xff0000)
                embed.set_thumbnail(url=result['thumbnail_url'])
                embed.set_author(name=result['channel_title'], url=result['channel_url'],
                    icon_url=result['channel_thumbnail_url'])
                await self.send_message(message.channel, embed=embed)
            else:
                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, server, user, message, filter, **kwargs):
        try:
            result = self.commands.find_clip(filter)
        except CommandError as e:
            await self.send_message(message.channel, 'Sorry {0}, {1}'.format(message.author.mention, e))
        else:
            await self.send_message(message.channel, '{title}\n{url}'.format(**result))