dinislage/CouchPotatoServer
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Subliminal update

Browse Source
This commit is contained in:
Ruud
2012-04-08 17:35:39 +02:00
parent 5321037f08
commit 148e1940ac
85 changed files with 4375 additions and 3420 deletions
Download Patch File Download Diff File Expand all files Collapse all files
libs/subliminal/core.py
+187 -313
View File
@@ -1,346 +1,173 @@
# -*- coding: utf-8 -*-
# Copyright 2011-2012 Antoine Bertin <diaoulael@gmail.com>
#
# Subliminal - Subtitles, faster than your thoughts
# Copyright (c) 2011 Antoine Bertin <diaoulael@gmail.com>
# This file is part of subliminal.
#
# This file is part of Subliminal.
#
# Subliminal is free software; you can redistribute it and/or modify it under
# the terms of the Lesser GNU General Public License as published by
# 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,
# 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
# Lesser GNU General Public License for more details.
# GNU Lesser General Public License for more details.
#
# You should have received a copy of the Lesser GNU General Public License
# along with this program. If not, see <http://www.gnu.org/licenses/>.
#
__all__ = ['PLUGINS', 'API_PLUGINS', 'IDLE', 'RUNNING', 'PAUSED', 'Subliminal', 'PluginWorker', 'matching_confidence',
'LANGUAGE_INDEX', 'PLUGIN_INDEX', 'PLUGIN_CONFIDENCE', 'MATCHING_CONFIDENCE']
# 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 collections import defaultdict
from exceptions import InvalidLanguageError, PluginError, BadStateError, \
WrongTaskError, DownloadFailedError
from itertools import groupby
from languages import list_languages
from utils import NullHandler
from tasks import Task, DownloadTask, ListTask, StopTask
import Queue
import guessit
import logging
import os
import plugins
import subtitles
import threading
import utils
import videos
# init logger
logger = logging.getLogger('subliminal')
logger.addHandler(NullHandler())
# const
PLUGINS = ['OpenSubtitles', 'BierDopje', 'TheSubDB', 'SubsWiki', 'Subtitulos']
API_PLUGINS = filter(lambda p: getattr(plugins, p).api_based, PLUGINS)
IDLE, RUNNING, PAUSED = range(3)
LANGUAGE_INDEX, PLUGIN_INDEX, PLUGIN_CONFIDENCE, MATCHING_CONFIDENCE = range(4)
class Subliminal(object):
"""Main Subliminal class"""
__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(__name__)
SERVICES = ['opensubtitles', 'bierdopje', 'subswiki', 'subtitulos', 'thesubdb']
LANGUAGE_INDEX, SERVICE_INDEX, SERVICE_CONFIDENCE, MATCHING_CONFIDENCE = range(4)
def __init__(self, cache_dir=None, workers=None, multi=False, force=False,
max_depth=None, filemode=None, sort_order=None, plugins=None, languages=None):
self.multi = multi
self.sort_order = sort_order or [LANGUAGE_INDEX, PLUGIN_INDEX, PLUGIN_CONFIDENCE, MATCHING_CONFIDENCE]
self.force = force
self.max_depth = max_depth or 3
self.taskQueue = Queue.PriorityQueue()
self.listResultQueue = Queue.Queue()
self.downloadResultQueue = Queue.Queue()
self.languages = languages or []
self.plugins = plugins or API_PLUGINS
self._workers = workers or 4
self.filemode = filemode
self.state = IDLE
self.cache_dir = cache_dir
try:
if cache_dir and not os.path.isdir(cache_dir):
os.makedirs(cache_dir)
logger.debug(u'Creating cache directory: %r' % cache_dir)
except:
self.cache_dir = None
logger.error(u'Failed to use the cache directory, continue without it')
def __enter__(self):
self.startWorkers()
return self
def create_list_tasks(paths, languages, services, force, multi, cache_dir, max_depth):
"""Create a list of :class:`~subliminal.tasks.ListTask` from one or more paths using the given criteria
def __exit__(self, *args):
self.stopWorkers(0)
: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
:return: the created tasks
:rtype: list of :class:`~subliminal.tasks.ListTask`
@property
def workers(self):
return self._workers
@workers.setter
def workers(self, value):
if self.state == RUNNING:
raise BadStateError(self.state, IDLE)
self._workers = value
@property
def languages(self):
"""Getter for languages"""
return self._languages
@languages.setter
def languages(self, languages):
"""Setter for languages"""
logger.debug(u'Setting languages to %r' % languages)
self._languages = []
for l in languages:
if l not in list_languages(1):
raise InvalidLanguageError(l)
if not l in self._languages:
self._languages.append(l)
@property
def plugins(self):
"""Getter for plugins"""
return self._plugins
@plugins.setter
def plugins(self, plugins):
"""Setter for plugins"""
logger.debug(u'Setting plugins to %r' % plugins)
self._plugins = []
for p in plugins:
if p not in PLUGINS:
raise PluginError(p)
if not p in self._plugins:
self._plugins.append(p)
def listSubtitles(self, entries, auto=False):
"""
Search subtitles within the plugins and return all found subtitles in a list of Subtitle object.
Attributes:
entries -- filepath or folderpath of video file or a list of that
auto -- automaticaly manage workers (default to False)"""
if auto:
if self.state != IDLE:
raise BadStateError(self.state, IDLE)
self.startWorkers()
if isinstance(entries, basestring):
entries = [entries]
config = utils.PluginConfig(self.multi, self.cache_dir, self.filemode)
scan_result = []
for e in entries:
if not isinstance(e, unicode):
logger.warning(u'Entry %r is not unicode' % e)
scan_result.extend(videos.scan(e))
task_count = 0
for video, subtitles in scan_result:
languages = set([s.language for s in subtitles if s.language])
wanted_languages = set(self._languages)
"""
scan_result = []
for p in paths:
scan_result.extend(scan(p, max_depth))
logger.debug(u'Found %d videos in %r with maximum depth %d' % (len(scan_result), paths, max_depth))
tasks = []
config = ServiceConfig(multi, cache_dir)
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:
wanted_languages = list_languages(1)
if not self.force and self.multi:
wanted_languages = set(wanted_languages) - languages
if not wanted_languages:
logger.debug(u'No need to list multi subtitles %r for %r because %r subtitles detected' % (self._languages, video.path, languages))
continue
if not self.force and not self.multi and None in [s.language for s in subtitles]:
logger.debug(u'No need to list single subtitles %r for %r because one detected' % (self._languages, video.path))
logger.debug(u'No need to list multi subtitles %r for %r because %r detected' % (languages, video, detected_languages))
continue
logger.debug(u'Listing subtitles %r for %r with %r' % (wanted_languages, video.path, self._plugins))
for plugin_name in self._plugins:
plugin = getattr(plugins, plugin_name)
to_list_languages = wanted_languages & plugin.availableLanguages()
if not to_list_languages:
logger.debug(u'Skipping %r: none of wanted languages %r available in %r for plugin %s' % (video.path, wanted_languages, plugin.availableLanguages(), plugin_name))
continue
if not plugin.isValidVideo(video):
logger.debug(u'Skipping %r: video %r is not part of supported videos %r for plugin %s' % (video.path, video, plugin.videos, plugin_name))
continue
self.taskQueue.put((5, ListTask(video, to_list_languages, plugin_name, config)))
task_count += 1
subtitles = []
for _ in range(task_count):
subtitles.extend(self.listResultQueue.get())
if auto:
self.stopWorkers()
return subtitles
def downloadSubtitles(self, entries, auto=False):
"""
Download subtitles using the plugins preferences and languages. Also use internal algorithm to find
the best match inside a plugin.
Attributes:
entries -- filepath or folderpath of video file or a list of that
auto -- automaticaly manage workers (default to False)"""
if auto:
if self.state != IDLE:
raise BadStateError(self.state, IDLE)
self.startWorkers()
by_video = self.groupByVideo(self.listSubtitles(entries, False))
# Define an order with LANGUAGE_INDEX first for multi sorting
order = self.sort_order
if self.multi:
order.insert(0, LANGUAGE_INDEX)
task_count = 0
for video, subtitles in by_video.iteritems():
ordered_subtitles = sorted(subtitles, key=lambda s: self.keySubtitles(s, video, order), reverse=True)
if not self.multi:
self.taskQueue.put((5, DownloadTask(video, list(ordered_subtitles))))
task_count += 1
if not force and not multi and None 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
service_languages = wanted_languages & service.available_languages()
if not service_languages:
logger.debug(u'Skipping %r: none of wanted languages %r available for service %s' % (video, wanted_languages, service_name))
continue
for _, by_language in groupby(ordered_subtitles, lambda s: s.language):
self.taskQueue.put((5, DownloadTask(video, list(by_language))))
task_count += 1
downloaded = []
for _ in range(task_count):
downloaded.extend(self.downloadResultQueue.get())
if auto:
self.stopWorkers()
return downloaded
def keySubtitles(self, subtitle, video, order):
"""Create a key to sort subtitle using preferences"""
key = ''
for sort_item in order:
if sort_item == LANGUAGE_INDEX:
key += '{0:03d}'.format(len(self._languages) - self._languages.index(subtitle.language) - 1)
elif sort_item == PLUGIN_INDEX:
key += '{0:02d}'.format(len(self._plugins) - self._plugins.index(subtitle.plugin) - 1)
elif sort_item == PLUGIN_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 groupByVideo(self, list_result):
'''Because list outputs a list of tuples from different plugins, we need to put them back
together under a single video key'''
result = defaultdict(list)
for video, subtitles in list_result:
result[video] += subtitles
return result
def startWorkers(self):
"""Create a pool of workers and start them"""
if self.state == RUNNING:
raise BadStateError(self.state, IDLE)
self.pool = []
for _ in range(self._workers):
worker = PluginWorker(self.taskQueue, self.listResultQueue, self.downloadResultQueue)
worker.start()
self.pool.append(worker)
logger.debug(u'Worker %s added to the pool' % worker.name)
self.state = RUNNING
def stopWorkers(self, priority=10):
"""Stop workers using a lowest priority stop signal and wait for them to terminate properly"""
for _ in range(self._workers):
self.taskQueue.put((priority, StopTask()))
for worker in self.pool:
worker.join()
self.state = IDLE
if not self.taskQueue.empty():
self.state = PAUSED
def pauseWorkers(self):
"""Pause workers using a highest priority stop signal and wait for them to terminate properly"""
self.stopWorkers(0)
def addTask(self, task):
"""Add a task with default priority"""
if not isinstance(task, Task) or isinstance(task, StopTask):
raise WrongTaskError()
self.taskQueue.put((5, task))
if not service.is_valid_video(video):
logger.debug(u'Skipping %r: not part of supported videos %r for service %s' % (video, service.videos, service_name))
continue
task = ListTask(video, service_languages, service_name, config)
logger.debug(u'Created task %r' % task)
tasks.append(task)
return tasks
class PluginWorker(threading.Thread):
"""Threaded plugin worker"""
def __init__(self, taskQueue, listResultQueue, downloadResultQueue):
threading.Thread.__init__(self)
self.taskQueue = taskQueue
self.listResultQueue = listResultQueue
self.downloadResultQueue = downloadResultQueue
self.logger = logging.getLogger('subliminal.worker')
self.plugins = {}
def create_download_tasks(subtitles_by_video, multi):
"""Create a list of :class:`~subliminal.tasks.DownloadTask` from a list results grouped by video
def run(self):
while True:
task = self.taskQueue.get()[1]
if isinstance(task, StopTask):
self.logger.debug(u'Poison pill received in thread %s' % self.name)
self.taskQueue.task_done()
:param subtitles_by_video: :class:`~subliminal.tasks.ListTask` results grouped by video and sorted
:type subtitles_by_video: dict of :class:`~subliminal.videos.Video` => [:class:`~subliminal.subtitles.Subtitle`]
:param order: preferred order for subtitles sorting
:type list: list of :data:`LANGUAGE_INDEX`, :data:`SERVICE_INDEX`, :data:`SERVICE_CONFIDENCE`, :data:`MATCHING_CONFIDENCE`
: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: 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` or :class:`~subliminal.subtitles.Subtitle`
"""
if services is None:
services = {}
logger.info(u'Consuming %r' % task)
result = None
if isinstance(task, ListTask):
if task.service not in services:
mod = __import__('services.' + task.service, globals=globals(), locals=locals(), fromlist=['Service'], level=-1)
services[task.service] = mod.Service(task.config)
services[task.service].init()
subtitles = services[task.service].list(task.video, task.languages)
result = subtitles
elif isinstance(task, DownloadTask):
for subtitle in task.subtitles:
if subtitle.service not in services:
mod = __import__('services.' + subtitle.service, globals=globals(), locals=locals(), fromlist=['Service'], level=-1)
services[subtitle.service] = mod.Service()
services[subtitle.service].init()
try:
services[subtitle.service].download(subtitle)
result = subtitle
break
result = []
try:
if isinstance(task, ListTask):
if task.plugin not in self.plugins: # init the plugin
self.plugins[task.plugin] = getattr(plugins, task.plugin)()
self.plugins[task.plugin].init()
# Retrieve the plugin list subtitles and return [(video, [subtitle])]
plugin = self.plugins[task.plugin]
plugin.config = task.config
subtitles = plugin.list(task.video, task.languages)
result = [(task.video, subtitles)]
elif isinstance(task, DownloadTask):
# Attempt to download one subtitle from the given list
for subtitle in task.subtitles:
if subtitle.plugin not in self.plugins: # init the plugin
self.plugins[subtitle.plugin] = getattr(plugins, subtitle.plugin)()
self.plugins[subtitle.plugin].init()
plugin = self.plugins[subtitle.plugin]
try:
result = [plugin.download(subtitle)]
break
except DownloadFailedError: # try the next one
self.logger.warning(u'Could not download subtitle %r, trying next' % subtitle)
continue
if not result:
self.logger.error(u'No subtitles could be downloaded for video %r' % task.video.path or task.video.release)
except:
self.logger.error(u'Exception raised in worker %s' % self.name, exc_info=True)
finally:
# Put the result in the correct queue
if isinstance(task, ListTask):
self.listResultQueue.put(result)
elif isinstance(task, DownloadTask):
self.downloadResultQueue.put(result)
self.taskQueue.task_done()
self.terminate()
self.logger.debug(u'Thread %s terminated' % self.name)
def terminate(self):
"""Terminate instanciated plugins"""
for plugin_name, plugin in self.plugins.iteritems():
try:
plugin.terminate()
except:
self.logger.error(u'Exception raised when terminating plugin %s' % plugin_name, exc_info=True)
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 confidence that the subtitle matches the video.
Returns a float between 0 and 1. 1 being the perfect match.'''
"""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 = utils.get_keywords(video.guess)
subtitle_keywords = utils.get_keywords(guess) | subtitle.keywords
video_keywords = get_keywords(video.guess)
subtitle_keywords = get_keywords(guess) | subtitle.keywords
replacement = {'keywords': len(video_keywords & subtitle_keywords)}
if isinstance(video, videos.Episode):
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))
@@ -351,7 +178,7 @@ def matching_confidence(video, subtitle):
replacement['season'] = 1
if 'episodeNumber' in guess and guess['episodeNumber'] == video.episode:
replacement['episode'] = 1
elif isinstance(video, videos.Movie):
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))
@@ -364,3 +191,50 @@ def matching_confidence(video, subtitle):
return 0
confidence = float(int(matching_format.format(**replacement), 2)) / float(int(best, 2))
return confidence
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)
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
return result