Move common libs to libs/common

libs/common/subliminal/refiners/__init__.py

@@ -0,0 +1,12 @@
+"""
+Refiners enrich a :class:`~subliminal.video.Video` object by adding information to it.
+
+A refiner is a simple function:
+
+.. py:function:: refine(video, **kwargs)
+
+    :param video: the video to refine.
+    :type video: :class:`~subliminal.video.Video`
+    :param \*\*kwargs: additional parameters for refiners.
+
+"""

libs/common/subliminal/refiners/metadata.py

@@ -0,0 +1,99 @@
+# -*- coding: utf-8 -*-
+import logging
+import os
+
+from babelfish import Error as BabelfishError, Language
+from enzyme import MKV
+
+logger = logging.getLogger(__name__)
+
+
+def refine(video, embedded_subtitles=True, **kwargs):
+    """Refine a video by searching its metadata.
+
+    Several :class:`~subliminal.video.Video` attributes can be found:
+
+      * :attr:`~subliminal.video.Video.resolution`
+      * :attr:`~subliminal.video.Video.video_codec`
+      * :attr:`~subliminal.video.Video.audio_codec`
+      * :attr:`~subliminal.video.Video.subtitle_languages`
+
+    :param bool embedded_subtitles: search for embedded subtitles.
+
+    """
+    # skip non existing videos
+    if not video.exists:
+        return
+
+    # check extensions
+    extension = os.path.splitext(video.name)[1]
+    if extension == '.mkv':
+        with open(video.name, 'rb') as f:
+            mkv = MKV(f)
+
+        # main video track
+        if mkv.video_tracks:
+            video_track = mkv.video_tracks[0]
+
+            # resolution
+            if video_track.height in (480, 720, 1080):
+                if video_track.interlaced:
+                    video.resolution = '%di' % video_track.height
+                else:
+                    video.resolution = '%dp' % video_track.height
+                logger.debug('Found resolution %s', video.resolution)
+
+            # video codec
+            if video_track.codec_id == 'V_MPEG4/ISO/AVC':
+                video.video_codec = 'h264'
+                logger.debug('Found video_codec %s', video.video_codec)
+            elif video_track.codec_id == 'V_MPEG4/ISO/SP':
+                video.video_codec = 'DivX'
+                logger.debug('Found video_codec %s', video.video_codec)
+            elif video_track.codec_id == 'V_MPEG4/ISO/ASP':
+                video.video_codec = 'XviD'
+                logger.debug('Found video_codec %s', video.video_codec)
+        else:
+            logger.warning('MKV has no video track')
+
+        # main audio track
+        if mkv.audio_tracks:
+            audio_track = mkv.audio_tracks[0]
+            # audio codec
+            if audio_track.codec_id == 'A_AC3':
+                video.audio_codec = 'AC3'
+                logger.debug('Found audio_codec %s', video.audio_codec)
+            elif audio_track.codec_id == 'A_DTS':
+                video.audio_codec = 'DTS'
+                logger.debug('Found audio_codec %s', video.audio_codec)
+            elif audio_track.codec_id == 'A_AAC':
+                video.audio_codec = 'AAC'
+                logger.debug('Found audio_codec %s', video.audio_codec)
+        else:
+            logger.warning('MKV has no audio track')
+
+        # subtitle tracks
+        if mkv.subtitle_tracks:
+            if embedded_subtitles:
+                embedded_subtitle_languages = set()
+                for st in mkv.subtitle_tracks:
+                    if st.language:
+                        try:
+                            embedded_subtitle_languages.add(Language.fromalpha3b(st.language))
+                        except BabelfishError:
+                            logger.error('Embedded subtitle track language %r is not a valid language', st.language)
+                            embedded_subtitle_languages.add(Language('und'))
+                    elif st.name:
+                        try:
+                            embedded_subtitle_languages.add(Language.fromname(st.name))
+                        except BabelfishError:
+                            logger.debug('Embedded subtitle track name %r is not a valid language', st.name)
+                            embedded_subtitle_languages.add(Language('und'))
+                    else:
+                        embedded_subtitle_languages.add(Language('und'))
+                logger.debug('Found embedded subtitle %r', embedded_subtitle_languages)
+                video.subtitle_languages |= embedded_subtitle_languages
+        else:
+            logger.debug('MKV has no subtitle track')
+    else:
+        logger.debug('Unsupported video extension %s', extension)

libs/common/subliminal/refiners/omdb.py

@@ -0,0 +1,187 @@
+# -*- coding: utf-8 -*-
+import logging
+import operator
+
+import requests
+
+from .. import __short_version__
+from ..cache import REFINER_EXPIRATION_TIME, region
+from ..video import Episode, Movie
+from ..utils import sanitize
+
+logger = logging.getLogger(__name__)
+
+
+class OMDBClient(object):
+    base_url = 'http://www.omdbapi.com'
+
+    def __init__(self, version=1, session=None, headers=None, timeout=10):
+        #: Session for the requests
+        self.session = session or requests.Session()
+        self.session.timeout = timeout
+        self.session.headers.update(headers or {})
+        self.session.params['r'] = 'json'
+        self.session.params['v'] = version
+
+    def get(self, id=None, title=None, type=None, year=None, plot='short', tomatoes=False):
+        # build the params
+        params = {}
+        if id:
+            params['i'] = id
+        if title:
+            params['t'] = title
+        if not params:
+            raise ValueError('At least id or title is required')
+        params['type'] = type
+        params['y'] = year
+        params['plot'] = plot
+        params['tomatoes'] = tomatoes
+
+        # perform the request
+        r = self.session.get(self.base_url, params=params)
+        r.raise_for_status()
+
+        # get the response as json
+        j = r.json()
+
+        # check response status
+        if j['Response'] == 'False':
+            return None
+
+        return j
+
+    def search(self, title, type=None, year=None, page=1):
+        # build the params
+        params = {'s': title, 'type': type, 'y': year, 'page': page}
+
+        # perform the request
+        r = self.session.get(self.base_url, params=params)
+        r.raise_for_status()
+
+        # get the response as json
+        j = r.json()
+
+        # check response status
+        if j['Response'] == 'False':
+            return None
+
+        return j
+
+
+omdb_client = OMDBClient(headers={'User-Agent': 'Subliminal/%s' % __short_version__})
+
+
+@region.cache_on_arguments(expiration_time=REFINER_EXPIRATION_TIME)
+def search(title, type, year):
+    results = omdb_client.search(title, type, year)
+    if not results:
+        return None
+
+    # fetch all paginated results
+    all_results = results['Search']
+    total_results = int(results['totalResults'])
+    page = 1
+    while total_results > page * 10:
+        page += 1
+        results = omdb_client.search(title, type, year, page=page)
+        all_results.extend(results['Search'])
+
+    return all_results
+
+
+def refine(video, **kwargs):
+    """Refine a video by searching `OMDb API <http://omdbapi.com/>`_.
+
+    Several :class:`~subliminal.video.Episode` attributes can be found:
+
+      * :attr:`~subliminal.video.Episode.series`
+      * :attr:`~subliminal.video.Episode.year`
+      * :attr:`~subliminal.video.Episode.series_imdb_id`
+
+    Similarly, for a :class:`~subliminal.video.Movie`:
+
+      * :attr:`~subliminal.video.Movie.title`
+      * :attr:`~subliminal.video.Movie.year`
+      * :attr:`~subliminal.video.Video.imdb_id`
+
+    """
+    if isinstance(video, Episode):
+        # exit if the information is complete
+        if video.series_imdb_id:
+            logger.debug('No need to search')
+            return
+
+        # search the series
+        results = search(video.series, 'series', video.year)
+        if not results:
+            logger.warning('No results for series')
+            return
+        logger.debug('Found %d results', len(results))
+
+        # filter the results
+        results = [r for r in results if sanitize(r['Title']) == sanitize(video.series)]
+        if not results:
+            logger.warning('No matching series found')
+            return
+
+        # process the results
+        found = False
+        for result in sorted(results, key=operator.itemgetter('Year')):
+            if video.original_series and video.year is None:
+                logger.debug('Found result for original series without year')
+                found = True
+                break
+            if video.year == int(result['Year'].split(u'\u2013')[0]):
+                logger.debug('Found result with matching year')
+                found = True
+                break
+
+        if not found:
+            logger.warning('No matching series found')
+            return
+
+        # add series information
+        logger.debug('Found series %r', result)
+        video.series = result['Title']
+        video.year = int(result['Year'].split(u'\u2013')[0])
+        video.series_imdb_id = result['imdbID']
+
+    elif isinstance(video, Movie):
+        # exit if the information is complete
+        if video.imdb_id:
+            return
+
+        # search the movie
+        results = search(video.title, 'movie', video.year)
+        if not results:
+            logger.warning('No results')
+            return
+        logger.debug('Found %d results', len(results))
+
+        # filter the results
+        results = [r for r in results if sanitize(r['Title']) == sanitize(video.title)]
+        if not results:
+            logger.warning('No matching movie found')
+            return
+
+        # process the results
+        found = False
+        for result in results:
+            if video.year is None:
+                logger.debug('Found result for movie without year')
+                found = True
+                break
+            if video.year == int(result['Year']):
+                logger.debug('Found result with matching year')
+                found = True
+                break
+
+        if not found:
+            logger.warning('No matching movie found')
+            return
+
+        # add movie information
+        logger.debug('Found movie %r', result)
+        video.title = result['Title']
+        video.year = int(result['Year'].split(u'\u2013')[0])
+        video.imdb_id = result['imdbID']

libs/common/subliminal/refiners/tvdb.py

@@ -0,0 +1,350 @@
+# -*- coding: utf-8 -*-
+from datetime import datetime, timedelta
+from functools import wraps
+import logging
+import re
+
+import requests
+
+from .. import __short_version__
+from ..cache import REFINER_EXPIRATION_TIME, region
+from ..utils import sanitize
+from ..video import Episode
+
+logger = logging.getLogger(__name__)
+
+series_re = re.compile(r'^(?P<series>.*?)(?: \((?:(?P<year>\d{4})|(?P<country>[A-Z]{2}))\))?$')
+
+
+def requires_auth(func):
+    """Decorator for :class:`TVDBClient` methods that require authentication"""
+    @wraps(func)
+    def wrapper(self, *args, **kwargs):
+        if self.token is None or self.token_expired:
+            self.login()
+        elif self.token_needs_refresh:
+            self.refresh_token()
+        return func(self, *args, **kwargs)
+    return wrapper
+
+
+class TVDBClient(object):
+    """TVDB REST API Client
+
+    :param str apikey: API key to use.
+    :param str username: username to use.
+    :param str password: password to use.
+    :param str language: language of the responses.
+    :param session: session object to use.
+    :type session: :class:`requests.sessions.Session` or compatible.
+    :param dict headers: additional headers.
+    :param int timeout: timeout for the requests.
+
+    """
+    #: Base URL of the API
+    base_url = 'https://api.thetvdb.com'
+
+    #: Token lifespan
+    token_lifespan = timedelta(hours=1)
+
+    #: Minimum token age before a :meth:`refresh_token` is triggered
+    refresh_token_every = timedelta(minutes=30)
+
+    def __init__(self, apikey=None, username=None, password=None, language='en', session=None, headers=None,
+                 timeout=10):
+        #: API key
+        self.apikey = apikey
+
+        #: Username
+        self.username = username
+
+        #: Password
+        self.password = password
+
+        #: Last token acquisition date
+        self.token_date = datetime.utcnow() - self.token_lifespan
+
+        #: Session for the requests
+        self.session = session or requests.Session()
+        self.session.timeout = timeout
+        self.session.headers.update(headers or {})
+        self.session.headers['Content-Type'] = 'application/json'
+        self.session.headers['Accept-Language'] = language
+
+    @property
+    def language(self):
+        return self.session.headers['Accept-Language']
+
+    @language.setter
+    def language(self, value):
+        self.session.headers['Accept-Language'] = value
+
+    @property
+    def token(self):
+        if 'Authorization' not in self.session.headers:
+            return None
+        return self.session.headers['Authorization'][7:]
+
+    @property
+    def token_expired(self):
+        return datetime.utcnow() - self.token_date > self.token_lifespan
+
+    @property
+    def token_needs_refresh(self):
+        return datetime.utcnow() - self.token_date > self.refresh_token_every
+
+    def login(self):
+        """Login"""
+        # perform the request
+        data = {'apikey': self.apikey, 'username': self.username, 'password': self.password}
+        r = self.session.post(self.base_url + '/login', json=data)
+        r.raise_for_status()
+
+        # set the Authorization header
+        self.session.headers['Authorization'] = 'Bearer ' + r.json()['token']
+
+        # update token_date
+        self.token_date = datetime.utcnow()
+
+    def refresh_token(self):
+        """Refresh token"""
+        # perform the request
+        r = self.session.get(self.base_url + '/refresh_token')
+        r.raise_for_status()
+
+        # set the Authorization header
+        self.session.headers['Authorization'] = 'Bearer ' + r.json()['token']
+
+        # update token_date
+        self.token_date = datetime.utcnow()
+
+    @requires_auth
+    def search_series(self, name=None, imdb_id=None, zap2it_id=None):
+        """Search series"""
+        # perform the request
+        params = {'name': name, 'imdbId': imdb_id, 'zap2itId': zap2it_id}
+        r = self.session.get(self.base_url + '/search/series', params=params)
+        if r.status_code == 404:
+            return None
+        r.raise_for_status()
+
+        return r.json()['data']
+
+    @requires_auth
+    def get_series(self, id):
+        """Get series"""
+        # perform the request
+        r = self.session.get(self.base_url + '/series/{}'.format(id))
+        if r.status_code == 404:
+            return None
+        r.raise_for_status()
+
+        return r.json()['data']
+
+    @requires_auth
+    def get_series_actors(self, id):
+        """Get series actors"""
+        # perform the request
+        r = self.session.get(self.base_url + '/series/{}/actors'.format(id))
+        if r.status_code == 404:
+            return None
+        r.raise_for_status()
+
+        return r.json()['data']
+
+    @requires_auth
+    def get_series_episodes(self, id, page=1):
+        """Get series episodes"""
+        # perform the request
+        params = {'page': page}
+        r = self.session.get(self.base_url + '/series/{}/episodes'.format(id), params=params)
+        if r.status_code == 404:
+            return None
+        r.raise_for_status()
+
+        return r.json()
+
+    @requires_auth
+    def query_series_episodes(self, id, absolute_number=None, aired_season=None, aired_episode=None, dvd_season=None,
+                              dvd_episode=None, imdb_id=None, page=1):
+        """Query series episodes"""
+        # perform the request
+        params = {'absoluteNumber': absolute_number, 'airedSeason': aired_season, 'airedEpisode': aired_episode,
+                  'dvdSeason': dvd_season, 'dvdEpisode': dvd_episode, 'imdbId': imdb_id, 'page': page}
+        r = self.session.get(self.base_url + '/series/{}/episodes/query'.format(id), params=params)
+        if r.status_code == 404:
+            return None
+        r.raise_for_status()
+
+        return r.json()
+
+    @requires_auth
+    def get_episode(self, id):
+        """Get episode"""
+        # perform the request
+        r = self.session.get(self.base_url + '/episodes/{}'.format(id))
+        if r.status_code == 404:
+            return None
+        r.raise_for_status()
+
+        return r.json()['data']
+
+
+#: Configured instance of :class:`TVDBClient`
+tvdb_client = TVDBClient('5EC930FB90DA1ADA', headers={'User-Agent': 'Subliminal/%s' % __short_version__})
+
+
+@region.cache_on_arguments(expiration_time=REFINER_EXPIRATION_TIME)
+def search_series(name):
+    """Search series.
+
+    :param str name: name of the series.
+    :return: the search results.
+    :rtype: list
+
+    """
+    return tvdb_client.search_series(name)
+
+
+@region.cache_on_arguments(expiration_time=REFINER_EXPIRATION_TIME)
+def get_series(id):
+    """Get series.
+
+    :param int id: id of the series.
+    :return: the series data.
+    :rtype: dict
+
+    """
+    return tvdb_client.get_series(id)
+
+
+@region.cache_on_arguments(expiration_time=REFINER_EXPIRATION_TIME)
+def get_series_episode(series_id, season, episode):
+    """Get an episode of a series.
+
+    :param int series_id: id of the series.
+    :param int season: season number of the episode.
+    :param int episode: episode number of the episode.
+    :return: the episode data.
+    :rtype: dict
+
+    """
+    result = tvdb_client.query_series_episodes(series_id, aired_season=season, aired_episode=episode)
+    if result:
+        return tvdb_client.get_episode(result['data'][0]['id'])
+
+
+def refine(video, **kwargs):
+    """Refine a video by searching `TheTVDB <http://thetvdb.com/>`_.
+
+    .. note::
+
+        This refiner only work for instances of :class:`~subliminal.video.Episode`.
+
+    Several attributes can be found:
+
+      * :attr:`~subliminal.video.Episode.series`
+      * :attr:`~subliminal.video.Episode.year`
+      * :attr:`~subliminal.video.Episode.series_imdb_id`
+      * :attr:`~subliminal.video.Episode.series_tvdb_id`
+      * :attr:`~subliminal.video.Episode.title`
+      * :attr:`~subliminal.video.Video.imdb_id`
+      * :attr:`~subliminal.video.Episode.tvdb_id`
+
+    """
+    # only deal with Episode videos
+    if not isinstance(video, Episode):
+        logger.error('Cannot refine episodes')
+        return
+
+    # exit if the information is complete
+    if video.series_tvdb_id and video.tvdb_id:
+        logger.debug('No need to search')
+        return
+
+    # search the series
+    logger.info('Searching series %r', video.series)
+    results = search_series(video.series.lower())
+    if not results:
+        logger.warning('No results for series')
+        return
+    logger.debug('Found %d results', len(results))
+
+    # search for exact matches
+    matching_results = []
+    for result in results:
+        matching_result = {}
+
+        # use seriesName and aliases
+        series_names = [result['seriesName']]
+        series_names.extend(result['aliases'])
+
+        # parse the original series as series + year or country
+        original_match = series_re.match(result['seriesName']).groupdict()
+
+        # parse series year
+        series_year = None
+        if result['firstAired']:
+            series_year = datetime.strptime(result['firstAired'], '%Y-%m-%d').year
+
+        # discard mismatches on year
+        if video.year and series_year and video.year != series_year:
+            logger.debug('Discarding series %r mismatch on year %d', result['seriesName'], series_year)
+            continue
+
+        # iterate over series names
+        for series_name in series_names:
+            # parse as series and year
+            series, year, country = series_re.match(series_name).groups()
+            if year:
+                year = int(year)
+
+            # discard mismatches on year
+            if year and (video.original_series or video.year != year):
+                logger.debug('Discarding series name %r mismatch on year %d', series, year)
+                continue
+
+            # match on sanitized series name
+            if sanitize(series) == sanitize(video.series):
+                logger.debug('Found exact match on series %r', series_name)
+                matching_result['match'] = {'series': original_match['series'], 'year': series_year,
+                                            'original_series': original_match['year'] is None}
+                break
+
+        # add the result on match
+        if matching_result:
+            matching_result['data'] = result
+            matching_results.append(matching_result)
+
+    # exit if we don't have exactly 1 matching result
+    if not matching_results:
+        logger.error('No matching series found')
+        return
+    if len(matching_results) > 1:
+        logger.error('Multiple matches found')
+        return
+
+    # get the series
+    matching_result = matching_results[0]
+    series = get_series(matching_result['data']['id'])
+
+    # add series information
+    logger.debug('Found series %r', series)
+    video.series = matching_result['match']['series']
+    video.year = matching_result['match']['year']
+    video.original_series = matching_result['match']['original_series']
+    video.series_tvdb_id = series['id']
+    video.series_imdb_id = series['imdbId'] or None
+
+    # get the episode
+    logger.info('Getting series episode %dx%d', video.season, video.episode)
+    episode = get_series_episode(video.series_tvdb_id, video.season, video.episode)
+    if not episode:
+        logger.warning('No results for episode')
+        return
+
+    # add episode information
+    logger.debug('Found episode %r', episode)
+    video.tvdb_id = episode['id']
+    video.title = episode['episodeName'] or None
+    video.imdb_id = episode['imdbId'] or None