mirror of
https://github.com/moparisthebest/SickRage
synced 2024-11-18 07:15:13 -05:00
276 lines
12 KiB
Python
276 lines
12 KiB
Python
|
# -*- coding: utf-8 -*-
|
||
|
|
# Copyright 2011-2012 Antoine Bertin <diaoulael@gmail.com>
|
||
|
|
#
|
||
|
|
# This file is part of subliminal.
|
||
|
|
#
|
||
|
|
# subliminal is free software; you can redistribute it and/or modify it under
|
||
|
|
# the terms of the GNU Lesser General Public License as published by
|
||
|
|
# the Free Software Foundation; either version 3 of the License, or
|
||
|
|
# (at your option) any later version.
|
||
|
|
#
|
||
|
|
# subliminal is distributed in the hope that it will be useful,
|
||
|
|
# but WITHOUT ANY WARRANTY; without even the implied warranty of
|
||
|
|
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
|
||
|
|
# GNU Lesser General Public License for more details.
|
||
|
|
#
|
||
|
|
# You should have received a copy of the GNU Lesser General Public License
|
||
|
|
# along with subliminal. If not, see <http://www.gnu.org/licenses/>.
|
||
|
|
from .exceptions import DownloadFailedError
|
||
|
|
from .services import ServiceConfig
|
||
|
|
from .tasks import DownloadTask, ListTask
|
||
|
|
from .utils import get_keywords
|
||
|
|
from .videos import Episode, Movie, scan
|
||
|
|
from .language import Language
|
||
|
|
from collections import defaultdict
|
||
|
|
from itertools import groupby
|
||
|
|
import bs4
|
||
|
|
import guessit
|
||
|
|
import logging
|
||
|
|
|
||
|
|
|
||
|
|
__all__ = ['SERVICES', 'LANGUAGE_INDEX', 'SERVICE_INDEX', 'SERVICE_CONFIDENCE', 'MATCHING_CONFIDENCE',
|
||
|
|
'create_list_tasks', 'create_download_tasks', 'consume_task', 'matching_confidence',
|
||
|
|
'key_subtitles', 'group_by_video']
|
||
|
|
logger = logging.getLogger("subliminal")
|
||
|
|
SERVICES = ['opensubtitles', 'subswiki', 'subtitulos', 'thesubdb', 'addic7ed', 'tvsubtitles', 'itasa', 'usub']
|
||
|
|
LANGUAGE_INDEX, SERVICE_INDEX, SERVICE_CONFIDENCE, MATCHING_CONFIDENCE = range(4)
|
||
|
|
|
||
|
|
|
||
|
|
def create_list_tasks(paths, languages, services, force, multi, cache_dir, max_depth, scan_filter):
|
||
|
|
"""Create a list of :class:`~subliminal.tasks.ListTask` from one or more paths using the given criteria
|
||
|
|
|
||
|
|
:param paths: path(s) to video file or folder
|
||
|
|
:type paths: string or list
|
||
|
|
:param set languages: languages to search for
|
||
|
|
:param list services: services to use for the search
|
||
|
|
:param bool force: force searching for subtitles even if some are detected
|
||
|
|
:param bool multi: search multiple languages for the same video
|
||
|
|
:param string cache_dir: path to the cache directory to use
|
||
|
|
:param int max_depth: maximum depth for scanning entries
|
||
|
|
:param function scan_filter: filter function that takes a path as argument and returns a boolean indicating whether it has to be filtered out (``True``) or not (``False``)
|
||
|
|
:return: the created tasks
|
||
|
|
:rtype: list of :class:`~subliminal.tasks.ListTask`
|
||
|
|
|
||
|
|
"""
|
||
|
|
scan_result = []
|
||
|
|
for p in paths:
|
||
|
|
scan_result.extend(scan(p, max_depth, scan_filter))
|
||
|
|
logger.debug(u'Found %d videos in %r with maximum depth %d' % (len(scan_result), paths, max_depth))
|
||
|
|
tasks = []
|
||
|
|
config = ServiceConfig(multi, cache_dir)
|
||
|
|
services = filter_services(services)
|
||
|
|
for video, detected_subtitles in scan_result:
|
||
|
|
detected_languages = set(s.language for s in detected_subtitles)
|
||
|
|
wanted_languages = languages.copy()
|
||
|
|
if not force and multi:
|
||
|
|
wanted_languages -= detected_languages
|
||
|
|
if not wanted_languages:
|
||
|
|
logger.debug(u'No need to list multi subtitles %r for %r because %r detected' % (languages, video, detected_languages))
|
||
|
|
continue
|
||
|
|
if not force and not multi and Language('Undetermined') in detected_languages:
|
||
|
|
logger.debug(u'No need to list single subtitles %r for %r because one detected' % (languages, video))
|
||
|
|
continue
|
||
|
|
logger.debug(u'Listing subtitles %r for %r with services %r' % (wanted_languages, video, services))
|
||
|
|
for service_name in services:
|
||
|
|
mod = __import__('services.' + service_name, globals=globals(), locals=locals(), fromlist=['Service'], level=-1)
|
||
|
|
service = mod.Service
|
||
|
|
if not service.check_validity(video, wanted_languages):
|
||
|
|
continue
|
||
|
|
task = ListTask(video, wanted_languages & service.languages, service_name, config)
|
||
|
|
logger.debug(u'Created task %r' % task)
|
||
|
|
tasks.append(task)
|
||
|
|
return tasks
|
||
|
|
|
||
|
|
|
||
|
|
def create_download_tasks(subtitles_by_video, languages, multi):
|
||
|
|
"""Create a list of :class:`~subliminal.tasks.DownloadTask` from a list results grouped by video
|
||
|
|
|
||
|
|
:param subtitles_by_video: :class:`~subliminal.tasks.ListTask` results with ordered subtitles
|
||
|
|
:type subtitles_by_video: dict of :class:`~subliminal.videos.Video` => [:class:`~subliminal.subtitles.Subtitle`]
|
||
|
|
:param languages: languages in preferred order
|
||
|
|
:type languages: :class:`~subliminal.language.language_list`
|
||
|
|
:param bool multi: download multiple languages for the same video
|
||
|
|
:return: the created tasks
|
||
|
|
:rtype: list of :class:`~subliminal.tasks.DownloadTask`
|
||
|
|
|
||
|
|
"""
|
||
|
|
tasks = []
|
||
|
|
for video, subtitles in subtitles_by_video.iteritems():
|
||
|
|
if not subtitles:
|
||
|
|
continue
|
||
|
|
if not multi:
|
||
|
|
task = DownloadTask(video, list(subtitles))
|
||
|
|
logger.debug(u'Created task %r' % task)
|
||
|
|
tasks.append(task)
|
||
|
|
continue
|
||
|
|
for _, by_language in groupby(subtitles, lambda s: languages.index(s.language)):
|
||
|
|
task = DownloadTask(video, list(by_language))
|
||
|
|
logger.debug(u'Created task %r' % task)
|
||
|
|
tasks.append(task)
|
||
|
|
return tasks
|
||
|
|
|
||
|
|
|
||
|
|
def consume_task(task, services=None):
|
||
|
|
"""Consume a task. If the ``services`` parameter is given, the function will attempt
|
||
|
|
to get the service from it. In case the service is not in ``services``, it will be initialized
|
||
|
|
and put in ``services``
|
||
|
|
|
||
|
|
:param task: task to consume
|
||
|
|
:type task: :class:`~subliminal.tasks.ListTask` or :class:`~subliminal.tasks.DownloadTask`
|
||
|
|
:param dict services: mapping between the service name and an instance of this service
|
||
|
|
:return: the result of the task
|
||
|
|
:rtype: list of :class:`~subliminal.subtitles.ResultSubtitle`
|
||
|
|
|
||
|
|
"""
|
||
|
|
if services is None:
|
||
|
|
services = {}
|
||
|
|
logger.info(u'Consuming %r' % task)
|
||
|
|
result = None
|
||
|
|
if isinstance(task, ListTask):
|
||
|
|
service = get_service(services, task.service, config=task.config)
|
||
|
|
result = service.list(task.video, task.languages)
|
||
|
|
elif isinstance(task, DownloadTask):
|
||
|
|
for subtitle in task.subtitles:
|
||
|
|
service = get_service(services, subtitle.service)
|
||
|
|
try:
|
||
|
|
service.download(subtitle)
|
||
|
|
result = [subtitle]
|
||
|
|
break
|
||
|
|
except DownloadFailedError:
|
||
|
|
logger.warning(u'Could not download subtitle %r, trying next' % subtitle)
|
||
|
|
continue
|
||
|
|
if result is None:
|
||
|
|
logger.error(u'No subtitles could be downloaded for video %r' % task.video)
|
||
|
|
return result
|
||
|
|
|
||
|
|
|
||
|
|
def matching_confidence(video, subtitle):
|
||
|
|
"""Compute the probability (confidence) that the subtitle matches the video
|
||
|
|
|
||
|
|
:param video: video to match
|
||
|
|
:type video: :class:`~subliminal.videos.Video`
|
||
|
|
:param subtitle: subtitle to match
|
||
|
|
:type subtitle: :class:`~subliminal.subtitles.Subtitle`
|
||
|
|
:return: the matching probability
|
||
|
|
:rtype: float
|
||
|
|
|
||
|
|
"""
|
||
|
|
guess = guessit.guess_file_info(subtitle.release, 'autodetect')
|
||
|
|
video_keywords = get_keywords(video.guess)
|
||
|
|
subtitle_keywords = get_keywords(guess) | subtitle.keywords
|
||
|
|
logger.debug(u'Video keywords %r - Subtitle keywords %r' % (video_keywords, subtitle_keywords))
|
||
|
|
replacement = {'keywords': len(video_keywords & subtitle_keywords)}
|
||
|
|
if isinstance(video, Episode):
|
||
|
|
replacement.update({'series': 0, 'season': 0, 'episode': 0})
|
||
|
|
matching_format = '{series:b}{season:b}{episode:b}{keywords:03b}'
|
||
|
|
best = matching_format.format(series=1, season=1, episode=1, keywords=len(video_keywords))
|
||
|
|
if guess['type'] in ['episode', 'episodesubtitle']:
|
||
|
|
if 'series' in guess and guess['series'].lower() == video.series.lower():
|
||
|
|
replacement['series'] = 1
|
||
|
|
if 'season' in guess and guess['season'] == video.season:
|
||
|
|
replacement['season'] = 1
|
||
|
|
if 'episodeNumber' in guess and guess['episodeNumber'] == video.episode:
|
||
|
|
replacement['episode'] = 1
|
||
|
|
elif isinstance(video, Movie):
|
||
|
|
replacement.update({'title': 0, 'year': 0})
|
||
|
|
matching_format = '{title:b}{year:b}{keywords:03b}'
|
||
|
|
best = matching_format.format(title=1, year=1, keywords=len(video_keywords))
|
||
|
|
if guess['type'] in ['movie', 'moviesubtitle']:
|
||
|
|
if 'title' in guess and guess['title'].lower() == video.title.lower():
|
||
|
|
replacement['title'] = 1
|
||
|
|
if 'year' in guess and guess['year'] == video.year:
|
||
|
|
replacement['year'] = 1
|
||
|
|
else:
|
||
|
|
logger.debug(u'Not able to compute confidence for %r' % video)
|
||
|
|
return 0.0
|
||
|
|
logger.debug(u'Found %r' % replacement)
|
||
|
|
confidence = float(int(matching_format.format(**replacement), 2)) / float(int(best, 2))
|
||
|
|
logger.info(u'Computed confidence %.4f for %r and %r' % (confidence, video, subtitle))
|
||
|
|
return confidence
|
||
|
|
|
||
|
|
|
||
|
|
def get_service(services, service_name, config=None):
|
||
|
|
"""Get a service from its name in the service dict with the specified config.
|
||
|
|
If the service does not exist in the service dict, it is created and added to the dict.
|
||
|
|
|
||
|
|
:param dict services: dict where to get existing services or put created ones
|
||
|
|
:param string service_name: name of the service to get
|
||
|
|
:param config: config to use for the service
|
||
|
|
:type config: :class:`~subliminal.services.ServiceConfig` or None
|
||
|
|
:return: the corresponding service
|
||
|
|
:rtype: :class:`~subliminal.services.ServiceBase`
|
||
|
|
|
||
|
|
"""
|
||
|
|
if service_name not in services:
|
||
|
|
mod = __import__('services.' + service_name, globals=globals(), locals=locals(), fromlist=['Service'], level=-1)
|
||
|
|
services[service_name] = mod.Service()
|
||
|
|
services[service_name].init()
|
||
|
|
services[service_name].config = config
|
||
|
|
return services[service_name]
|
||
|
|
|
||
|
|
|
||
|
|
def key_subtitles(subtitle, video, languages, services, order):
|
||
|
|
"""Create a key to sort subtitle using the given order
|
||
|
|
|
||
|
|
:param subtitle: subtitle to sort
|
||
|
|
:type subtitle: :class:`~subliminal.subtitles.ResultSubtitle`
|
||
|
|
:param video: video to match
|
||
|
|
:type video: :class:`~subliminal.videos.Video`
|
||
|
|
:param list languages: languages in preferred order
|
||
|
|
:param list services: services in preferred order
|
||
|
|
:param order: preferred order for subtitles sorting
|
||
|
|
:type list: list of :data:`LANGUAGE_INDEX`, :data:`SERVICE_INDEX`, :data:`SERVICE_CONFIDENCE`, :data:`MATCHING_CONFIDENCE`
|
||
|
|
:return: a key ready to use for subtitles sorting
|
||
|
|
:rtype: int
|
||
|
|
|
||
|
|
"""
|
||
|
|
key = ''
|
||
|
|
for sort_item in order:
|
||
|
|
if sort_item == LANGUAGE_INDEX:
|
||
|
|
key += '{0:03d}'.format(len(languages) - languages.index(subtitle.language) - 1)
|
||
|
|
key += '{0:01d}'.format(subtitle.language == languages[languages.index(subtitle.language)])
|
||
|
|
elif sort_item == SERVICE_INDEX:
|
||
|
|
key += '{0:02d}'.format(len(services) - services.index(subtitle.service) - 1)
|
||
|
|
elif sort_item == SERVICE_CONFIDENCE:
|
||
|
|
key += '{0:04d}'.format(int(subtitle.confidence * 1000))
|
||
|
|
elif sort_item == MATCHING_CONFIDENCE:
|
||
|
|
confidence = 0
|
||
|
|
if subtitle.release:
|
||
|
|
confidence = matching_confidence(video, subtitle)
|
||
|
|
key += '{0:04d}'.format(int(confidence * 1000))
|
||
|
|
return int(key)
|
||
|
|
|
||
|
|
|
||
|
|
def group_by_video(list_results):
|
||
|
|
"""Group the results of :class:`ListTasks <subliminal.tasks.ListTask>` into a
|
||
|
|
dictionary of :class:`~subliminal.videos.Video` => :class:`~subliminal.subtitles.Subtitle`
|
||
|
|
|
||
|
|
:param list_results:
|
||
|
|
:type list_results: list of result of :class:`~subliminal.tasks.ListTask`
|
||
|
|
:return: subtitles grouped by videos
|
||
|
|
:rtype: dict of :class:`~subliminal.videos.Video` => [:class:`~subliminal.subtitles.Subtitle`]
|
||
|
|
|
||
|
|
"""
|
||
|
|
result = defaultdict(list)
|
||
|
|
for video, subtitles in list_results:
|
||
|
|
result[video] += subtitles or []
|
||
|
|
return result
|
||
|
|
|
||
|
|
|
||
|
|
def filter_services(services):
|
||
|
|
"""Filter out services that are not available because of a missing feature
|
||
|
|
|
||
|
|
:param list services: service names to filter
|
||
|
|
:return: a copy of the initial list of service names without unavailable ones
|
||
|
|
:rtype: list
|
||
|
|
|
||
|
|
"""
|
||
|
|
filtered_services = services[:]
|
||
|
|
for service_name in services:
|
||
|
|
mod = __import__('services.' + service_name, globals=globals(), locals=locals(), fromlist=['Service'], level=-1)
|
||
|
|
service = mod.Service
|
||
|
|
if service.required_features is not None and bs4.builder_registry.lookup(*service.required_features) is None:
|
||
|
|
logger.warning(u'Service %s not available: none of available features could be used. One of %r required' % (service_name, service.required_features))
|
||
|
|
filtered_services.remove(service_name)
|
||
|
|
return filtered_services
|