Fixed maximum number of requests per minute for assrt provider. #1953

pull/1974/head
morpheus65535 2 years ago
parent af70cf1fc9
commit 167cb5dd98

@ -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()

Loading…
Cancel
Save