diff options
| -rw-r--r-- | libs/subliminal_patch/providers/assrt.py | 52 |
1 files changed, 36 insertions, 16 deletions
diff --git a/libs/subliminal_patch/providers/assrt.py b/libs/subliminal_patch/providers/assrt.py index 6a101bba9..9173e7691 100644 --- a/libs/subliminal_patch/providers/assrt.py +++ b/libs/subliminal_patch/providers/assrt.py @@ -1,6 +1,4 @@ # -*- coding: utf-8 -*- -from __future__ import absolute_import -import json import logging import os import re @@ -9,6 +7,7 @@ from babelfish import language_converters from guessit import guessit from requests import Session from time import sleep +from math import ceil from subliminal import Movie, Episode, ProviderError, __short_version__ from subliminal.exceptions import AuthenticationError, ConfigurationError, DownloadLimitExceeded, ProviderError @@ -24,34 +23,42 @@ language_converters.register('assrt = subliminal_patch.converters.assrt:AssrtCon server_url = 'https://api.assrt.net/v1' supported_languages = list(language_converters['assrt'].to_assrt.keys()) + +def get_request_delay(max_request_per_minute): + return ceil(60 / max_request_per_minute) + + def language_contains(subset, superset): if subset.alpha3 != superset.alpha3: return False - if superset.country != None and subset.country != superset.country: + if superset.country is not None and subset.country != superset.country: return False - if superset.script != None and subset.script != superset.script: + if superset.script is not None and subset.script != superset.script: return False return True + def search_language_in_list(lang, langlist): - for l in langlist: - if language_contains(lang, l): - return l + for language in langlist: + if language_contains(lang, language): + return language return None + class AssrtSubtitle(Subtitle): """Assrt Sbutitle.""" provider_name = 'assrt' guessit_options = { - 'allowed_languages': [ l[0] for l in supported_languages ], + 'allowed_languages': [lang[0] for lang in supported_languages], # 'allowed_countries': [ l[1] for l in supported_languages if len(l) > 1 ], 'enforce_list': True } - def __init__(self, language, subtitle_id, video_name, session, token): + def __init__(self, language, subtitle_id, video_name, session, token, max_request_per_minute): super(AssrtSubtitle, self).__init__(language) self.session = session self.token = token + self.max_request_per_minute = max_request_per_minute self.subtitle_id = subtitle_id self.video_name = video_name self.release_info = video_name @@ -63,7 +70,7 @@ class AssrtSubtitle(Subtitle): return self._detail params = {'token': self.token, 'id': self.id} logger.info('Get subtitle detail: GET /sub/detail %r', params) - sleep(3) + sleep(get_request_delay(self.max_request_per_minute)) r = self.session.get(server_url + '/sub/detail', params=params, timeout=10) r.raise_for_status() @@ -86,7 +93,7 @@ class AssrtSubtitle(Subtitle): # second pass: keyword matching codes = language_converters['assrt'].codes for f in files: - langs = set([ Language.fromassrt(k) for k in codes if k in f['f'] ]) + langs = set([Language.fromassrt(k) for k in codes if k in f['f']]) if self.language in langs: self._defail = f return f @@ -110,17 +117,29 @@ class AssrtSubtitle(Subtitle): class AssrtProvider(Provider): """Assrt Provider.""" - languages = {Language(*l) for l in supported_languages} + languages = {Language(*lang) for lang in supported_languages} video_types = (Episode, Movie) def __init__(self, token=None): if not token: raise ConfigurationError('Token must be specified') self.token = token + self.session = Session() + self.default_max_request_per_minute = 20 + self.max_request_per_minute = None def initialize(self): - self.session = Session() self.session.headers = {'User-Agent': os.environ.get("SZ_USER_AGENT", "Sub-Zero/2")} + res = self.session.get(server_url + '/user/quota', params={'token': self.token}, timeout=10) + res.raise_for_status() + try: + result = res.json() + except: + self.max_request_per_minute = self.default_max_request_per_minute + else: + if 'user' in result: + if 'quota' in result['user']: + self.max_request_per_minute = result['user']['quota'] def terminate(self): self.session.close() @@ -148,7 +167,7 @@ class AssrtProvider(Provider): params = {'token': self.token, 'q': query, 'is_file': 1} logger.debug('Searching subtitles: GET /sub/search %r', params) - sleep(3) + sleep(get_request_delay(self.max_request_per_minute)) res = self.session.get(server_url + '/sub/search', params=params, timeout=10) res.raise_for_status() result = res.json() @@ -169,7 +188,8 @@ class AssrtProvider(Provider): language = Language.fromassrt(match.group('code')) output_language = search_language_in_list(language, languages) if output_language: - subtitles.append(AssrtSubtitle(output_language, sub['id'], sub['videoname'], self.session, self.token)) + subtitles.append(AssrtSubtitle(output_language, sub['id'], sub['videoname'], self.session, + self.token, self.max_request_per_minute)) except: pass @@ -179,7 +199,7 @@ class AssrtProvider(Provider): return self.query(languages, video) def download_subtitle(self, subtitle): - sleep(3) + sleep(get_request_delay(self.max_request_per_minute)) r = self.session.get(subtitle.download_link, timeout=10) r.raise_for_status() |