diff --git a/bot.py b/bot.py index bd4a1fb..d477cf2 100644 --- a/bot.py +++ b/bot.py @@ -5,11 +5,14 @@ import configparser import json import logging import pathlib +import re import discord import persistqueue import tweepy +from youtube import YouTube, YouTubeError + TIMEOUT = 1 @@ -18,6 +21,8 @@ TWITTER_STATUS_URL = 'https://twitter.com/i/status/' TWITTER_ICON_URL = 'https://abs.twimg.com/icons/apple-touch-icon-192x192.png' TWITTER_COLOR = 0x1da1f2 +YOUTUBE_COLOR = 0xff0000 + logger = logging.getLogger('discord') logger.addHandler(logging.StreamHandler()) @@ -36,8 +41,18 @@ channel_id = config.getint('Discord', 'channel_id') consumer_key = config.get('Twitter', 'consumer_key') consumer_secret = config.get('Twitter', 'consumer_secret') +yt_api_key = config.get('YouTube', 'api_key') +yt_channel_id = config.get('YouTube', 'channel_id') + class Bot(discord.Client): + def __init__(self): + self.commands = ( + re.compile(r'^(?P!)(?Preaction)\s+(?P")?(?P.+)(?(q)")$'), + ) + self.youtube = YouTube(yt_api_key, yt_channel_id) + super().__init__() + def get_api(self, user_id): handler = tweepy.OAuthHandler(consumer_key, consumer_secret) with open(pathlib.Path(storage_path, 'tokens.json'), 'r') as f: @@ -88,6 +103,35 @@ class Bot(discord.Client): for tweet in tweets: await channel.send(embed=self.make_embed(tweet)) + async def on_message(self, message): + for pattern in self.commands: + m = pattern.match(message.content) + if not m: + continue + d = m.groupdict() + command = d.pop('command') + action = getattr(self, 'perform_{0}'.format(command)) + await action(message, **d) + + async def perform_reaction(self, message, query, **kwargs): + try: + result = self.youtube.search(query) + except YouTubeError as e: + await message.channel.send('Sorry {0}, {1}'.format(message.author.mention, str(e))) + return + if not result: + await message.channel.send('Sorry {0}, nothing found'.format(message.author.mention)) + return + if result.get('kind') == 'playlist': + embed = discord.Embed(title=result.get('title'), url=result.get('link'), + description=result.get('description'), color=YOUTUBE_COLOR) + embed.set_thumbnail(url=result.get('thumbnail_url')) + embed.set_author(name=self.youtube.channel.get('title'), url=self.youtube.channel.get('link'), + icon_url=self.youtube.channel.get('thumbnail_url')) + await message.channel.send(embed=embed) + else: + await message.channel.send('{title}\n{link}'.format(**result)) + def main(): bot = Bot() diff --git a/youtube.py b/youtube.py new file mode 100644 index 0000000..0b743ac --- /dev/null +++ b/youtube.py @@ -0,0 +1,120 @@ +import html + +import fuzzywuzzy.fuzz +import fuzzywuzzy.process +import googleapiclient +import googleapiclient.discovery + + +BASE_URL = 'https://www.youtube.com' + + +class YouTube(object): + def __init__(self, api_key, channel_id): + self.client = googleapiclient.discovery.build('youtube', 'v3', developerKey=api_key) + self.channel = self.get_channel(channel_id) + + def get_thumbnail_url(self, thumbnails): + for key in ('high', 'medium', 'default'): + if key in thumbnails: + return thumbnails[key].get('url') + + def process_items(self, items): + result = [] + for item in items: + id = item.get('id', '') + kind = item.get('kind', 'youtube#').split('youtube#')[1] + if kind == 'searchResult': + id = item.get('id', {}).get('videoId', '') + kind = item.get('id', {}).get('kind', 'youtube#').split('youtube#')[1] + if kind == 'playlist': + link = '{0}/view_play_list?p={1}'.format(BASE_URL, id) + else: + link = '{0}/watch?v={1}'.format(BASE_URL, id) + result.append(dict( + kind=kind, + link=link, + title=html.unescape(item.get('snippet', {}).get('title', '')), + description=html.unescape(item.get('snippet', {}).get('description', '')), + thumbnail_url=self.get_thumbnail_url(item.get('snippet', {}).get('thumbnails', {})) + )) + return result + + def get_channel(self, channel_id): + r = self.client.channels().list(id=channel_id, maxResults=1, part='id,snippet').execute() + channel = r.get('items', [{}]).pop() + return dict( + id=channel.get('id', ''), + link='{0}/c/{1}'.format(BASE_URL, channel.get('snippet', {}).get('customUrl', '')), + title=html.unescape(channel.get('snippet', {}).get('title', '')), + thumbnail_url=self.get_thumbnail_url(channel.get('snippet', {}).get('thumbnails', {})) + ) + + def get_playlists(self): + token = '' + result = [] + while True: + r = self.client.playlists().list( + channelId=self.channel.get('id'), + maxResults=50, + part='id,snippet', + pageToken=token + ).execute() + result.extend(self.process_items(r.get('items', []))) + token = r.get('nextPageToken') + if not token: + break + return result + + def search_videos(self, query, limit): + count = limit + token = '' + result = [] + while True: + r = self.client.search().list( + channelId=self.channel.get('id'), + q=query, + safeSearch='none', + type='video', + maxResults=min(count, 50), + part='id,snippet', + pageToken=token + ).execute() + result.extend(self.process_items(r.get('items', []))) + count -= r.get('pageInfo', {}).get('resultsPerPage', 0) + if count <= 0: + break + token = r.get('nextPageToken') + if not token: + break + return result + + def search(self, query): + try: + results = self.get_playlists() + results.extend(self.search_videos(query, limit=5)) + except googleapiclient.errors.HttpError as e: + raise YouTubeError(str(e)) + if not results: + return None + tokens = [t for t in query.split('|') if not t.strip().startswith('-')] or [''] + matches = [] + for token in tokens: + titles = {i: r.get('title') for i, r in enumerate(results)} + descriptions = {i: r.get('description') for i, r in enumerate(results)} + match = fuzzywuzzy.process.extractOne(token, titles, + scorer=fuzzywuzzy.fuzz.token_set_ratio, score_cutoff=25) + if match: + matches.append(match) + match = fuzzywuzzy.process.extractOne(token, descriptions, + scorer=fuzzywuzzy.fuzz.token_set_ratio, score_cutoff=25) + if match: + matches.append(match) + if not matches: + return None + _, _, i = sorted(matches, key=lambda m: m[1], reverse=True)[0] + return results[i] + + +class YouTubeError(Exception): + pass