Add retry attempts to Requests downloader

This commit is contained in:
rlaphoenix 2024-01-09 12:08:49 +00:00
parent fa3cee11b7
commit 552a0f13a5
1 changed files with 32 additions and 18 deletions

View File

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