1
0
mirror of https://github.com/moparisthebest/SickRage synced 2024-10-31 23:45:02 -04:00
SickRage/lib/subliminal/services/__init__.py
echel0n c945726f05 Fixed subliminal issues.
Added ignore/required words option to bet set individually for each show.
Fixed issue with global ignore words not properly matching against releases.
Fixed issue with
2014-04-23 22:18:16 -07:00

260 lines
9.1 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 ..cache import Cache
from ..exceptions import DownloadFailedError, ServiceError
from ..language import language_set, Language
from ..subtitles import EXTENSIONS
import logging
import os
import requests
import threading
import zipfile
__all__ = ['ServiceBase', 'ServiceConfig']
logger = logging.getLogger("subliminal")
class ServiceBase(object):
"""Service base class
:param config: service configuration
:type config: :class:`ServiceConfig`
"""
#: URL to the service server
server_url = ''
#: User Agent for any HTTP-based requests
user_agent = 'subliminal v0.6'
#: Whether based on an API or not
api_based = False
#: Timeout for web requests
timeout = 5
#: :class:`~subliminal.language.language_set` of available languages
languages = language_set()
#: Map between language objects and language codes used in the service
language_map = {}
#: Default attribute of a :class:`~subliminal.language.Language` to get with :meth:`get_code`
language_code = 'alpha2'
#: Accepted video classes (:class:`~subliminal.videos.Episode`, :class:`~subliminal.videos.Movie`, :class:`~subliminal.videos.UnknownVideo`)
videos = []
#: Whether the video has to exist or not
require_video = False
#: List of required features for BeautifulSoup
required_features = None
def __init__(self, config=None):
self.config = config or ServiceConfig()
self.session = None
def __enter__(self):
self.init()
return self
def __exit__(self, *args):
self.terminate()
def init(self):
"""Initialize connection"""
logger.debug(u'Initializing %s' % self.__class__.__name__)
self.session = requests.session()
self.session.headers.update({'User-Agent': self.user_agent})
def init_cache(self):
"""Initialize cache, make sure it is loaded from disk"""
if not self.config or not self.config.cache:
raise ServiceError('Cache directory is required')
self.config.cache.load(self.__class__.__name__)
def save_cache(self):
self.config.cache.save(self.__class__.__name__)
def clear_cache(self):
self.config.cache.clear(self.__class__.__name__)
def cache_for(self, func, args, result):
return self.config.cache.cache_for(self.__class__.__name__, func, args, result)
def cached_value(self, func, args):
return self.config.cache.cached_value(self.__class__.__name__, func, args)
def terminate(self):
"""Terminate connection"""
logger.debug(u'Terminating %s' % self.__class__.__name__)
def get_code(self, language):
"""Get the service code for a :class:`~subliminal.language.Language`
It uses the :data:`language_map` and if there's no match, falls back
on the :data:`language_code` attribute of the given :class:`~subliminal.language.Language`
"""
if language in self.language_map:
return self.language_map[language]
if self.language_code is None:
raise ValueError('%r has no matching code' % language)
return getattr(language, self.language_code)
def get_language(self, code):
"""Get a :class:`~subliminal.language.Language` from a service code
It uses the :data:`language_map` and if there's no match, uses the
given code as ``language`` parameter for the :class:`~subliminal.language.Language`
constructor
.. note::
A warning is emitted if the generated :class:`~subliminal.language.Language`
is "Undetermined"
"""
if code in self.language_map:
return self.language_map[code]
language = Language(code, strict=False)
if language == Language('Undetermined'):
logger.warning(u'Code %s could not be identified as a language for %s' % (code, self.__class__.__name__))
return language
def query(self, *args):
"""Make the actual query"""
raise NotImplementedError()
def list(self, video, languages):
"""List subtitles
As a service writer, you can either override this method or implement
:meth:`list_checked` instead to have the languages pre-filtered for you
"""
if not self.check_validity(video, languages):
return []
return self.list_checked(video, languages)
def list_checked(self, video, languages):
"""List subtitles without having to check parameters for validity"""
raise NotImplementedError()
def download(self, subtitle):
"""Download a subtitle"""
self.download_file(subtitle.link, subtitle.path)
return subtitle
@classmethod
def check_validity(cls, video, languages):
"""Check for video and languages validity in the Service
:param video: the video to check
:type video: :class:`~subliminal.videos.video`
:param languages: languages to check
:type languages: :class:`~subliminal.language.Language`
:rtype: bool
"""
languages = (languages & cls.languages) - language_set(['Undetermined'])
if not languages:
logger.debug(u'No language available for service %s' % cls.__name__.lower())
return False
if cls.require_video and not video.exists or not isinstance(video, tuple(cls.videos)):
logger.debug(u'%r is not valid for service %s' % (video, cls.__name__.lower()))
return False
return True
def download_file(self, url, filepath):
"""Attempt to download a file and remove it in case of failure
:param string url: URL to download
:param string filepath: destination path
"""
logger.info(u'Downloading %s in %s' % (url, filepath))
try:
r = self.session.get(url, timeout = 10, headers = {'Referer': url, 'User-Agent': self.user_agent})
with open(filepath, 'wb') as f:
f.write(r.content)
except Exception as e:
logger.error(u'Download failed: %s' % e)
if os.path.exists(filepath):
os.remove(filepath)
raise DownloadFailedError(str(e))
logger.debug(u'Download finished')
def download_zip_file(self, url, filepath):
"""Attempt to download a zip file and extract any subtitle file from it, if any.
This cleans up after itself if anything fails.
:param string url: URL of the zip file to download
:param string filepath: destination path for the subtitle
"""
logger.info(u'Downloading %s in %s' % (url, filepath))
try:
zippath = filepath + '.zip'
r = self.session.get(url, timeout = 10, headers = {'Referer': url, 'User-Agent': self.user_agent})
with open(zippath, 'wb') as f:
f.write(r.content)
if not zipfile.is_zipfile(zippath):
# TODO: could check if maybe we already have a text file and
# download it directly
raise DownloadFailedError('Downloaded file is not a zip file')
zipsub = zipfile.ZipFile(zippath)
for subfile in zipsub.namelist():
if os.path.splitext(subfile)[1] in EXTENSIONS:
with open(filepath, 'wb') as f:
f.write(zipsub.open(subfile).read())
break
else:
zipsub.close()
raise DownloadFailedError('No subtitles found in zip file')
zipsub.close()
os.remove(zippath)
except Exception as e:
logger.error(u'Download %s failed: %s' % (url, e))
if os.path.exists(zippath):
os.remove(zippath)
if os.path.exists(filepath):
os.remove(filepath)
raise DownloadFailedError(str(e))
logger.debug(u'Download finished')
class ServiceConfig(object):
"""Configuration for any :class:`Service`
:param bool multi: whether to download one subtitle per language or not
:param string cache_dir: cache directory
"""
def __init__(self, multi=False, cache_dir=None):
self.multi = multi
self.cache_dir = cache_dir
self.cache = None
if cache_dir is not None:
self.cache = Cache(cache_dir)
def __repr__(self):
return 'ServiceConfig(%r, %s)' % (self.multi, self.cache.cache_dir)