SickRage/lib/subliminal/api.py

# -*- coding: utf-8 -*-
from __future__ import unicode_literals
import collections
import io
import logging
import operator
import os.path
import babelfish
from .providers import ProviderPool
from .subtitle import get_subtitle_path


logger = logging.getLogger(__name__)


def list_subtitles(videos, languages, providers=None, provider_configs=None):
    """List subtitles for `videos` with the given `languages` using the specified `providers`

    :param videos: videos to list subtitles for
    :type videos: set of :class:`~subliminal.video.Video`
    :param languages: languages of subtitles to search for
    :type languages: set of :class:`babelfish.Language`
    :param providers: providers to use, if not all
    :type providers: list of string or None
    :param provider_configs: configuration for providers
    :type provider_configs: dict of provider name => provider constructor kwargs or None
    :return: found subtitles
    :rtype: dict of :class:`~subliminal.video.Video` => [:class:`~subliminal.subtitle.Subtitle`]

    """
    subtitles = collections.defaultdict(list)
    with ProviderPool(providers, provider_configs) as pp:
        for video in videos:
            logger.info('Listing subtitles for %r', video)
            video_subtitles = pp.list_subtitles(video, languages)
            logger.info('Found %d subtitles total', len(video_subtitles))
            subtitles[video].extend(video_subtitles)
    return subtitles


def download_subtitles(subtitles, provider_configs=None):
    """Download subtitles

    :param subtitles: subtitles to download
    :type subtitles: list of :class:`~subliminal.subtitle.Subtitle`
    :param provider_configs: configuration for providers
    :type provider_configs: dict of provider name => provider constructor kwargs or None

    """
    with ProviderPool(provider_configs=provider_configs) as pp:
        for subtitle in subtitles:
            logger.info('Downloading subtitle %r', subtitle)
            pp.download_subtitle(subtitle)


def download_best_subtitles(videos, languages, providers=None, provider_configs=None, min_score=0,
                            hearing_impaired=False, single=False):
    """Download the best subtitles for `videos` with the given `languages` using the specified `providers`

    :param videos: videos to download subtitles for
    :type videos: set of :class:`~subliminal.video.Video`
    :param languages: languages of subtitles to download
    :type languages: set of :class:`babelfish.Language`
    :param providers: providers to use for the search, if not all
    :type providers: list of string or None
    :param provider_configs: configuration for providers
    :type provider_configs: dict of provider name => provider constructor kwargs or None
    :param int min_score: minimum score for subtitles to download
    :param bool hearing_impaired: download hearing impaired subtitles
    :param bool single: do not download for videos with an undetermined subtitle language detected

    """
    downloaded_subtitles = collections.defaultdict(list)
    with ProviderPool(providers, provider_configs) as pp:
        for video in videos:
            # filter
            if single and babelfish.Language('und') in video.subtitle_languages:
                logger.debug('Skipping video %r: undetermined language found')
                continue

            # list
            logger.info('Listing subtitles for %r', video)
            video_subtitles = pp.list_subtitles(video, languages)
            logger.info('Found %d subtitles total', len(video_subtitles))

            # download
            downloaded_languages = set()
            for subtitle, score in sorted([(s, s.compute_score(video)) for s in video_subtitles],
                                          key=operator.itemgetter(1), reverse=True):
                if score < min_score:
                    logger.info('No subtitle with score >= %d', min_score)
                    break
                if subtitle.hearing_impaired != hearing_impaired:
                    logger.debug('Skipping subtitle: hearing impaired != %r', hearing_impaired)
                    continue
                if subtitle.language in downloaded_languages:
                    logger.debug('Skipping subtitle: %r already downloaded', subtitle.language)
                    continue
                logger.info('Downloading subtitle %r with score %d', subtitle, score)
                if pp.download_subtitle(subtitle):
                    downloaded_languages.add(subtitle.language)
                    downloaded_subtitles[video].append(subtitle)
                if single or downloaded_languages == languages:
                    logger.debug('All languages downloaded')
                    break
    return downloaded_subtitles


def save_subtitles(subtitles, single=False, directory=None, encoding=None):
    """Save subtitles on disk next to the video or in a specific folder if `folder_path` is specified

    :param bool single: download with .srt extension if ``True``, add language identifier otherwise
    :param directory: path to directory where to save the subtitles, if any
    :type directory: string or None
    :param encoding: encoding for the subtitles or ``None`` to use the original encoding
    :type encoding: string or None

    """
    for video, video_subtitles in subtitles.items():
        saved_languages = set()
        for video_subtitle in video_subtitles:
            if video_subtitle.content is None:
                logger.debug('Skipping subtitle %r: no content', video_subtitle)
                continue
            if video_subtitle.language in saved_languages:
                logger.debug('Skipping subtitle %r: language already saved', video_subtitle)
                continue
            subtitle_path = get_subtitle_path(video.name, None if single else video_subtitle.language)
            if directory is not None:
                subtitle_path = os.path.join(directory, os.path.split(subtitle_path)[1])
            logger.info('Saving %r to %r', video_subtitle, subtitle_path)
            if encoding is None:
                with io.open(subtitle_path, 'wb') as f:
                    f.write(video_subtitle.content)
            else:
                with io.open(subtitle_path, 'w', encoding=encoding) as f:
                    f.write(video_subtitle.text)
            saved_languages.add(video_subtitle.language)
            if single:
                break