dfd

ProviderCore/download.py

@@ -1,7 +1,7 @@
 from __future__ import annotations
 
 from pathlib import Path
-from typing import Optional
+from typing import Callable, Optional
 import sys
 
 import requests
@@ -22,13 +22,20 @@ def sanitize_filename(name: str, *, max_len: int = 150) -> str:
     return cleaned[:max_len]
 
 
-def download_file(url: str, output_path: Path, *, session: Optional[requests.Session] = None, timeout_s: float = 30.0) -> bool:
+def download_file(
+    url: str,
+    output_path: Path,
+    *,
+    session: Optional[requests.Session] = None,
+    timeout_s: float = 30.0,
+    progress_callback: Optional[Callable[[int, Optional[int], str], None]] = None,
+) -> bool:
     output_path = Path(output_path)
     output_path.parent.mkdir(parents=True, exist_ok=True)
 
     s = session or requests.Session()
 
-    bar = ProgressBar()
+    bar = ProgressBar() if progress_callback is None else None
     downloaded = 0
     total = None
 
@@ -41,9 +48,14 @@ def download_file(url: str, output_path: Path, *, session: Optional[requests.Ses
             except Exception:
                 total = None
 
+            label = str(output_path.name or "download")
+
             # Render once immediately so fast downloads still show something.
             try:
-                bar.update(downloaded=0, total=total, label=str(output_path.name or "download"), file=sys.stderr)
+                if progress_callback is not None:
+                    progress_callback(0, total, label)
+                elif bar is not None:
+                    bar.update(downloaded=0, total=total, label=label, file=sys.stderr)
             except Exception:
                 pass
 
@@ -53,18 +65,23 @@ def download_file(url: str, output_path: Path, *, session: Optional[requests.Ses
                         f.write(chunk)
                         downloaded += len(chunk)
                         try:
-                            bar.update(downloaded=downloaded, total=total, label=str(output_path.name or "download"), file=sys.stderr)
+                            if progress_callback is not None:
+                                progress_callback(downloaded, total, label)
+                            elif bar is not None:
+                                bar.update(downloaded=downloaded, total=total, label=label, file=sys.stderr)
                         except Exception:
                             pass
 
         try:
-            bar.finish()
+            if bar is not None:
+                bar.finish()
         except Exception:
             pass
         return output_path.exists() and output_path.stat().st_size > 0
     except Exception:
         try:
-            bar.finish()
+            if bar is not None:
+                bar.finish()
         except Exception:
             pass
         try: