Improve youtube matching

master
Nikola Forró 6 years ago
parent df4be019eb
commit e2ef6c5588

@ -4,6 +4,9 @@ WORKDIR /bot
COPY . . COPY . .
RUN touch settings.cfg RUN touch settings.cfg
RUN apk add --no-cache gcc libc-dev && \
rm -rf /var/cache/apk/*
RUN pip install --no-cache-dir --requirement requirements.txt RUN pip install --no-cache-dir --requirement requirements.txt
RUN addgroup -g 9999 lilia RUN addgroup -g 9999 lilia

@ -104,15 +104,8 @@ class Commands(object):
channel_ids = self.config['Youtube'].get('channel_ids').split(',') channel_ids = self.config['Youtube'].get('channel_ids').split(',')
yt = Youtube(api_key) yt = Youtube(api_key)
try: try:
for channel_id in channel_ids: result = yt.find_best_match(channel_ids, query)
results = yt.search(channel_id, query, playlists=True, limit=1) except YoutubeError:
if results:
break
results = yt.search(channel_id, query, playlists=False, limit=1)
if results:
break
result = results.pop(0)
except (YoutubeError, IndexError):
raise CommandError('couldn\'t find anything on Youtube') raise CommandError('couldn\'t find anything on Youtube')
else: else:
return result return result

@ -1,4 +1,5 @@
discord.py discord.py
fuzzywuzzy[speedup]
google-api-python-client google-api-python-client
irc irc
python-dateutil python-dateutil

@ -1,3 +1,5 @@
import fuzzywuzzy.fuzz
import fuzzywuzzy.process
import googleapiclient import googleapiclient
import googleapiclient.discovery import googleapiclient.discovery
@ -63,3 +65,25 @@ class Youtube(object):
return self._search(channel_id, query, playlists, limit) return self._search(channel_id, query, playlists, limit)
except googleapiclient.errors.HttpError as e: except googleapiclient.errors.HttpError as e:
raise YoutubeError('Failed to query Youtube API: {}'.format(e)) raise YoutubeError('Failed to query Youtube API: {}'.format(e))
def find_best_match(self, channel_ids, query):
results = []
for channel_id in channel_ids:
try:
results.extend(self._search(channel_id, query, playlists=True, limit=1))
results.extend(self._search(channel_id, query, playlists=False, limit=1))
except googleapiclient.errors.HttpError as e:
raise YoutubeError('Failed to query Youtube API: {}'.format(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['title'] for i, r in enumerate(results)}
descriptions = {i: r['description'] for i, r in enumerate(results)}
matches.append(fuzzywuzzy.process.extractOne(token, titles,
scorer=fuzzywuzzy.fuzz.token_sort_ratio))
matches.append(fuzzywuzzy.process.extractOne(token, descriptions,
scorer=fuzzywuzzy.fuzz.token_sort_ratio))
_, _, i = sorted(matches, key=lambda m: m[1], reverse=True)[0]
return results[i]

Loading…
Cancel
Save