From 552a0f13a58f9bbc88b509bc570a30d4a0e6c970 Mon Sep 17 00:00:00 2001
From: rlaphoenix <rlaphoenix@pm.me>
Date: Tue, 9 Jan 2024 12:08:49 +0000
Subject: [PATCH] Add retry attempts to Requests downloader

---
 devine/core/downloaders/requests.py | 50 ++++++++++++++++++-----------
 1 file changed, 32 insertions(+), 18 deletions(-)

diff --git a/devine/core/downloaders/requests.py b/devine/core/downloaders/requests.py
index 53c5606..891807d 100644
--- a/devine/core/downloaders/requests.py
+++ b/devine/core/downloaders/requests.py
@@ -7,6 +7,11 @@ from requests import Session
 from requests.cookies import RequestsCookieJar
 from rich import filesize
 
+from devine.core.constants import DOWNLOAD_CANCELLED
+
+MAX_ATTEMPTS = 5
+RETRY_WAIT = 2
+
 
 def requests(
     uri: Union[str, list[str]],
@@ -59,26 +64,35 @@ def requests(
 
     for url, out_path in uri:
         out_path.parent.mkdir(parents=True, exist_ok=True)
-        stream = session.get(url, stream=True)
-        with open(out_path, "wb") as f:
-            written = 0
-            for chunk in stream.iter_content(chunk_size=1024):
-                download_size = len(chunk)
-                f.write(chunk)
-                written += download_size
-                if progress:
-                    progress(advance=1)
+        attempts = 1
+        try:
+            stream = session.get(url, stream=True)
+            stream.raise_for_status()
+            with open(out_path, "wb") as f:
+                written = 0
+                for chunk in stream.iter_content(chunk_size=1024):
+                    download_size = len(chunk)
+                    f.write(chunk)
+                    written += download_size
+                    if progress:
+                        progress(advance=1)
 
-                    now = time.time()
-                    time_since = now - last_speed_refresh
+                        now = time.time()
+                        time_since = now - last_speed_refresh
 
-                    download_sizes.append(download_size)
-                    if time_since > 5 or download_size < 1024:
-                        data_size = sum(download_sizes)
-                        download_speed = data_size / (time_since or 1)
-                        progress(downloaded=f"{filesize.decimal(download_speed)}/s")
-                        last_speed_refresh = now
-                        download_sizes.clear()
+                        download_sizes.append(download_size)
+                        if time_since > 5 or download_size < 1024:
+                            data_size = sum(download_sizes)
+                            download_speed = data_size / (time_since or 1)
+                            progress(downloaded=f"{filesize.decimal(download_speed)}/s")
+                            last_speed_refresh = now
+                            download_sizes.clear()
+            break
+        except Exception as e:
+            if DOWNLOAD_CANCELLED.is_set() or attempts == MAX_ATTEMPTS:
+                raise e
+            time.sleep(RETRY_WAIT)
+            attempts += 1
 
     return 0