mirror of
https://github.com/moparisthebest/SickRage
synced 2024-11-14 05:15:12 -05:00
c945726f05
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
260 lines
9.1 KiB
Python
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)
|