bot/services/twitch.py

import dateutil.parser
import requests


class TwitchError(Exception):
    pass


class Twitch(object):
    def __init__(self, api_url, client_id):
        self.api_url = api_url
        self.client_id = client_id

    def _get_videos(self, channel_id):
        def request(offset, limit):
            url = '{0}/channels/{1}/videos'.format(self.api_url, channel_id)
            params = dict(client_id=self.client_id, offset=offset, limit=limit)
            r = requests.get(url, params=params)
            r.raise_for_status()
            return r.json()
        result = []
        data = request(0, 1)
        total = data.get('_total', 0)
        limit = 100
        for offset in range(0, total, limit):
            data = request(offset, limit)
            for vid in data.get('videos', []):
                result.append(dict(
                    id=int(vid['_id'].lstrip('v')),
                    title=vid['title'],
                    date=dateutil.parser.parse(vid['recorded_at']).date()))
        return result

    def _get_comments(self, video_id):
        def request(cursor):
            url = '{0}/videos/{1}/comments'.format(self.api_url, video_id)
            params = dict(client_id=self.client_id, cursor=cursor)
            r = requests.get(url, params=params)
            r.raise_for_status()
            return r.json()
        result = []
        cursor = ''
        while True:
            data = request(cursor)
            for comment in data.get('comments', []):
                result.append(comment['message']['body'])
            cursor = data.get('_next')
            if not cursor:
                break
        return result

    def _get_stream_info(self, channel_id):
        def request():
            url = '{0}/streams/{1}'.format(self.api_url, channel_id)
            params = dict(client_id=self.client_id)
            r = requests.get(url, params=params)
            r.raise_for_status()
            return r.json()
        data = request()
        if data['stream'] is None:
            return None
        return dict(
            id=int(data['stream']['_id']),
            game=data['stream']['game'],
            viewers=int(data['stream']['viewers']))

    def get_messages(self, channel_id, since):
        try:
            videos = self._get_videos(channel_id)
            result = []
            for video in [v for v in videos if v['date'] >= since]:
                comments = self._get_comments(video['id'])
                result.extend(comments)
        except requests.exceptions.HTTPError as e:
            raise TwitchError('Failed to retrieve VOD messages: {0}'.format(e))
        else:
            return result

    def get_stream_info(self, channel_id):
        try:
            return self._get_stream_info(channel_id)
        except requests.exceptions.HTTPError as e:
            raise TwitchError('Failed to get stream info: {0}'.format(e))