SickGear/lib/libtrakt/indexerapiinterface.py

import logging
import re
import time
from .exceptions import TraktShowNotFound, TraktException
from sickbeard.exceptions import ex
from trakt import TraktAPI


class ShowContainer(dict):
    """Simple dict that holds a series of Show instances
    """

    def __init__(self, **kwargs):
        super(ShowContainer, self).__init__(**kwargs)

        self._stack = []
        self._lastgc = time.time()

    def __setitem__(self, key, value):

        self._stack.append(key)

        # keep only the 100th latest results
        if time.time() - self._lastgc > 20:
            for o in self._stack[:-100]:
                del self[o]

            self._stack = self._stack[-100:]

            self._lastgc = time.time()

        super(ShowContainer, self).__setitem__(key, value)


def log():
    return logging.getLogger('trakt_api')


class TraktSearchTypes:
    text = 1
    trakt_id = 'trakt'
    tvdb_id = 'tvdb'
    imdb_id = 'imdb'
    tmdb_id = 'tmdb'
    tvrage_id = 'tvrage'
    all = [text, trakt_id, tvdb_id, imdb_id, tmdb_id, tvrage_id]

    def __init__(self):
        pass


class TraktResultTypes:
    show = 'show'
    episode = 'episode'
    movie = 'movie'
    person = 'person'
    list = 'list'
    all = [show, episode, movie, person, list]

    def __init__(self):
        pass


class TraktIndexer:
    # noinspection PyUnusedLocal
    # noinspection PyDefaultArgument
    def __init__(self, custom_ui=None, sleep_retry=None, search_type=TraktSearchTypes.text,
                 result_types=[TraktResultTypes.show], *args, **kwargs):

        self.config = {
            'apikey': '',
            'debug_enabled': False,
            'custom_ui': custom_ui,
            'proxy': None,
            'cache_enabled': False,
            'cache_location': '',
            'valid_languages': [],
            'langabbv_to_id': {},
            'language': 'en',
            'base_url': '',
            'search_type': search_type if search_type in TraktSearchTypes.all else TraktSearchTypes.text,
            'sleep_retry': sleep_retry,
            'result_types': result_types if isinstance(result_types, list) and all(x in TraktResultTypes.all for x in result_types) else [TraktResultTypes.show],
        }

        self.corrections = {}
        self.shows = ShowContainer()

    def _get_series(self, series):
        """This searches Trakt for the series name,
        If a custom_ui UI is configured, it uses this to select the correct
        series.
        """
        all_series = self.search(series)
        if not isinstance(all_series, list):
            all_series = [all_series]

        if 0 == len(all_series):
            log().debug('Series result returned zero')
            raise TraktShowNotFound('Show-name search returned zero results (cannot find show on TVDB)')

        if None is not self.config['custom_ui']:
            log().debug('Using custom UI %s' % (repr(self.config['custom_ui'])))
            custom_ui = self.config['custom_ui']
            ui = custom_ui(config=self.config)

            return ui.select_series(all_series)

        return all_series

    def __getitem__(self, key):
        """Handles trakt_instance['seriesname'] calls.
        The dict index should be the show id
        """
        if isinstance(key, tuple) and 2 == len(key):
            key = key[0]

        self.config['searchterm'] = key
        selected_series = self._get_series(key)
        if isinstance(selected_series, dict):
            selected_series = [selected_series]

        return selected_series

    def __repr__(self):
        return str(self.shows)

    def _clean_data(self, data):
        """Cleans up strings, lists, dicts returned

        Issues corrected:
        - Replaces & with &
        - Trailing whitespace
        """
        if isinstance(data, list):
            return [self._clean_data(d) for d in data]
        if isinstance(data, dict):
            return {k: self._clean_data(v) for k, v in data.iteritems()}
        return data if not isinstance(data, (str, unicode)) else data.strip().replace(u'&', u'&')

    @staticmethod
    def _dict_prevent_none(d, key, default):
        v = None
        if isinstance(d, dict):
            v = d.get(key, default)
        return (v, default)[None is v]

    def search(self, series):
        if TraktSearchTypes.text != self.config['search_type']:
            url = '/search/%s/%s?type=%s&extended=full&limit=100' % (self.config['search_type'], series,
                                                                     ','.join(self.config['result_types']))
        else:
            url = '/search/%s?query=%s&extended=full&limit=100' % (','.join(self.config['result_types']), series)
        filtered = []
        kwargs = {}
        if None is not self.config['sleep_retry']:
            kwargs['sleep_retry'] = self.config['sleep_retry']
        try:
            resp = TraktAPI().trakt_request(url, **kwargs)
            if len(resp):
                for d in resp:
                    if isinstance(d, dict) and 'type' in d and d['type'] in self.config['result_types']:
                        for k, v in d.iteritems():
                            d[k] = self._clean_data(v)
                        if 'show' in d and TraktResultTypes.show == d['type']:
                            d.update(d['show'])
                            del d['show']
                            d['seriesname'] = self._dict_prevent_none(d, 'title', '')
                            d['genres_list'] = d.get('genres', [])
                            d['genres'] = ', '.join(['%s' % v for v in d.get('genres', []) or [] if v])
                            d['firstaired'] = (d.get('first_aired') and
                                               re.sub(r'T.*$', '', str(d.get('first_aired'))) or d.get('year'))
                        filtered.append(d)
        except TraktException as e:
            log().debug('Could not connect to Trakt service: %s' % ex(e))

        return filtered