Merge pull request #2 from robxu9/main

Add timeouts/retries

humblebundle_downloader/download_library.py

@@ -8,6 +8,9 @@
 import requests
 import http.cookiejar
 
+from requests.adapters import HTTPAdapter
+from requests.packages.urllib3.util.retry import Retry
+
 logger = logging.getLogger(__name__)
 
 
@@ -21,6 +24,24 @@ def _clean_name(dirty_str):
     return "".join(clean).strip().rstrip('.')
 
 
+DEFAULT_TIMEOUT = 5  # seconds
+
+
+class TimeoutHTTPAdapter(HTTPAdapter):
+    def __init__(self, *args, **kwargs):
+        self.timeout = DEFAULT_TIMEOUT
+        if "timeout" in kwargs:
+            self.timeout = kwargs["timeout"]
+            del kwargs["timeout"]
+        super().__init__(*args, **kwargs)
+
+    def send(self, request, **kwargs):
+        timeout = kwargs.get("timeout")
+        if timeout is None:
+            kwargs["timeout"] = self.timeout
+        return super().send(request, **kwargs)
+
+
 class DownloadLibrary:
 
     def __init__(self, library_path, cookie_path=None, cookie_auth=None,
@@ -42,7 +63,13 @@ def __init__(self, library_path, cookie_path=None, cookie_auth=None,
         self.update = update
         self.content_types = ['web'] if content_types is None else list(map(str.lower, content_types))  # noqa: E501
 
+        retries = Retry(total=3, backoff_factor=1,
+                        status_forcelist=[429, 500, 502, 503, 504])
+        timeout_adapter = TimeoutHTTPAdapter(max_retries=retries)
+
         self.session = requests.Session()
+        self.session.mount('http://', timeout_adapter)
+        self.session.mount('https://', timeout_adapter)
         if cookie_path:
             try:
                 cookie_jar = http.cookiejar.MozillaCookieJar(cookie_path)
@@ -132,8 +159,10 @@ def _process_trove_product(self, title, product):
                     self.library_path, 'Humble Trove', title
                 )
                 # Create directory to save the files to
-                try: os.makedirs(product_folder)  # noqa: E701
-                except OSError: pass  # noqa: E701
+                try:
+                    os.makedirs(product_folder)  # noqa: E701
+                except OSError:
+                    pass  # noqa: E701
                 local_filename = os.path.join(
                     product_folder,
                     web_name,
@@ -241,8 +270,10 @@ def _process_product(self, order_id, bundle_title, product):
                 self.library_path, bundle_title, product_title
             )
             # Create directory to save the files to
-            try: os.makedirs(product_folder)  # noqa: E701
-            except OSError: pass  # noqa: E701
+            try:
+                os.makedirs(product_folder)  # noqa: E701
+            except OSError:
+                pass  # noqa: E701
 
             # Download each file type of a product
             for file_type in download_type['download_struct']:
@@ -334,8 +365,10 @@ def _process_download(self, open_r, cache_file_key, file_info,
                          .format(local_filename=local_filename))
 
             # Clean up broken downloaded file
-            try: os.remove(local_filename)  # noqa: E701
-            except OSError: pass  # noqa: E701
+            try:
+                os.remove(local_filename)  # noqa: E701
+            except OSError:
+                pass  # noqa: E701
 
             if type(e).__name__ == 'KeyboardInterrupt':
                 sys.exit()