add subliminal for subtitle download. #253

libs/subliminal/api.py

@@ -0,0 +1,140 @@
+# -*- 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