DL tracker fixed!!!! now make it nice!

packages/node/meshnet_node/downloader.py

@@ -122,6 +122,7 @@ def _fetch_source_file(
     file_url: str,
     dest: Path,
     timeout: float,
+    on_chunk=None,
 ) -> int:
     """Download one file; skip if *dest* already matches the remote size.
 
@@ -135,6 +136,8 @@ def _fetch_source_file(
         length = resp.getheader("Content-Length") if hasattr(resp, "getheader") else None
         expected = int(length) if isinstance(length, str) and length.isdigit() else None
         if expected is not None and dest.exists() and dest.stat().st_size == expected:
+            if on_chunk is not None:
+                on_chunk(expected)
             return expected  # complete from an earlier attempt/run
         received = 0
         with dest.open("wb") as out:
@@ -144,11 +147,63 @@ def _fetch_source_file(
                     break
                 out.write(chunk)
                 received += len(chunk)
+                if on_chunk is not None:
+                    on_chunk(len(chunk))
         if expected is not None and received != expected:
             raise OSError(f"short read: got {received} of {expected} bytes")
         return received
 
 
+class _SourceProgress:
+    """tqdm bar over the whole per-file download (total / speed / ETA),
+    falling back to plain per-file prints when tqdm is unavailable."""
+
+    def __init__(self, enabled: bool, label: str, total_bytes: int | None):
+        self._label = label
+        self._bar = None
+        self._enabled = enabled
+        if enabled and total_bytes:
+            try:
+                from tqdm import tqdm  # type: ignore[import]
+
+                self._bar = tqdm(
+                    total=total_bytes,
+                    unit="B",
+                    unit_scale=True,
+                    unit_divisor=1024,
+                    desc=f"Downloading ({label})",
+                )
+            except Exception:
+                self._bar = None
+
+    def add_bytes(self, n: int) -> None:
+        if self._bar is not None:
+            self._bar.update(n)
+
+    def rewind(self, n: int) -> None:
+        """Roll the bar back after a failed attempt so retries don't double-count."""
+        if self._bar is not None and n:
+            self._bar.update(-n)
+
+    def file_done(self, index: int, count: int, rel: str, size: int, reused: bool) -> None:
+        note = "already complete, " if reused else ""
+        line = f"  {self._label}: [{index}/{count}] {rel} ({note}{size / 1e9:.2f} GB)"
+        if self._bar is not None:
+            self._bar.set_postfix_str(f"{index}/{count} files", refresh=False)
+        elif self._enabled:
+            print(line, flush=True)
+
+    def message(self, text: str) -> None:
+        if self._bar is not None:
+            self._bar.write(text)
+        elif self._enabled:
+            print(text, flush=True)
+
+    def close(self) -> None:
+        if self._bar is not None:
+            self._bar.close()
+
+
 def _reuse_local_file(expected: int, dest: Path, final: Path) -> bool:
     """Reuse an already-complete copy of a file instead of re-downloading."""
     if dest.exists() and dest.stat().st_size == expected:
@@ -188,41 +243,46 @@ def _download_source_files(
     partial_dir = shard_dir.parent / f"{shard_dir.name}.partial"
     partial_dir.mkdir(parents=True, exist_ok=True)
     sep = "&" if "?" in url else "?"
-    for index, rel in enumerate(rel_files, start=1):
-        dest = partial_dir / rel
-        dest.parent.mkdir(parents=True, exist_ok=True)
-        known_size = sizes.get(rel)
-        if isinstance(known_size, int) and _reuse_local_file(known_size, dest, shard_dir / rel):
-            if progress:
-                print(
-                    f"  {label}: [{index}/{len(rel_files)}] {rel} "
-                    f"(already complete, {known_size / 1e9:.2f} GB)",
-                    flush=True,
-                )
-            continue
-        file_url = f"{url}{sep}{urllib.parse.urlencode({'file': rel})}"
-        for attempt in range(1, _FILE_RETRY_ATTEMPTS + 1):
-            try:
-                size = _fetch_source_file(file_url, dest, timeout)
-                if progress:
-                    print(
-                        f"  {label}: [{index}/{len(rel_files)}] {rel} ({size / 1e9:.2f} GB)",
-                        flush=True,
-                    )
-                break
-            except _TarOnlySource:
-                if progress:
-                    print(f"  {label}: no single-file support — using tar stream", flush=True)
-                return None
-            except Exception as exc:
-                if progress:
-                    print(
-                        f"  {label}: {rel} attempt {attempt}/{_FILE_RETRY_ATTEMPTS} failed: {exc!r}",
-                        flush=True,
-                    )
-                if attempt == _FILE_RETRY_ATTEMPTS:
+    known_sizes = [sizes.get(rel) for rel in rel_files]
+    total_bytes = sum(s for s in known_sizes if isinstance(s, int)) if all(
+        isinstance(s, int) for s in known_sizes
+    ) else None
+    tracker_bar = _SourceProgress(progress, label, total_bytes)
+    try:
+        for index, rel in enumerate(rel_files, start=1):
+            dest = partial_dir / rel
+            dest.parent.mkdir(parents=True, exist_ok=True)
+            known_size = sizes.get(rel)
+            if isinstance(known_size, int) and _reuse_local_file(known_size, dest, shard_dir / rel):
+                tracker_bar.add_bytes(known_size)
+                tracker_bar.file_done(index, len(rel_files), rel, known_size, reused=True)
+                continue
+            file_url = f"{url}{sep}{urllib.parse.urlencode({'file': rel})}"
+            for attempt in range(1, _FILE_RETRY_ATTEMPTS + 1):
+                counted = 0
+
+                def _on_chunk(n: int) -> None:
+                    nonlocal counted
+                    counted += n
+                    tracker_bar.add_bytes(n)
+
+                try:
+                    size = _fetch_source_file(file_url, dest, timeout, on_chunk=_on_chunk)
+                    tracker_bar.file_done(index, len(rel_files), rel, size, reused=False)
+                    break
+                except _TarOnlySource:
+                    tracker_bar.message(f"  {label}: no single-file support — using tar stream")
                     return None
-                time.sleep(1.0 * attempt)
+                except Exception as exc:
+                    tracker_bar.rewind(counted)
+                    tracker_bar.message(
+                        f"  {label}: {rel} attempt {attempt}/{_FILE_RETRY_ATTEMPTS} failed: {exc!r}"
+                    )
+                    if attempt == _FILE_RETRY_ATTEMPTS:
+                        return None
+                    time.sleep(1.0 * attempt)
+    finally:
+        tracker_bar.close()
     if shard_dir.exists():
         shutil.rmtree(shard_dir)
     shutil.move(str(partial_dir), str(shard_dir))