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

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