Keep the weights download progress output short in logs

CHANGELOG.md

@@ -16,6 +16,12 @@ All notable changes to SubCellSeg are listed here. The format follows
   (was 10%). Fine-tuning uses a learning rate of 5e-5 and stops after 50
   epochs without improvement; training from scratch keeps 1e-4 and 100.
 
+### Fixed
+- The progress line for the first-time weights download no longer floods
+  logs: in a terminal it redraws once per percent, and when output goes to a
+  file, CI or a notebook it writes one line per 10% (was about 8,000 redraws,
+  ~150 KB). ([#5](https://github.com/AllenCell/SubCellSeg/issues/5))
+
 ### Added
 - Benchmarks section in the README and `tools/benchmark/` scripts.
 - Contributing guide, code of conduct, security policy, citation file,

src/subcellseg/weights.py

@@ -32,14 +32,40 @@ def sha256(path: Path) -> str:
     return digest.hexdigest()
 
 
-def _report(blocks: int, block_size: int, total: int) -> None:
-    done = blocks * block_size / 2**20
-    if total > 0:
-        pct = min(100, int(done * 2**20 * 100 / total))
-        sys.stderr.write(f"\r  {done:.0f} / {total / 2**20:.0f} MB ({pct}%)")
-    else:
-        sys.stderr.write(f"\r  {done:.0f} MB")
-    sys.stderr.flush()
+class _Progress:
+    """``urlretrieve`` reporthook that writes only when the shown value changes.
+
+    urlretrieve calls the hook for every 8 KB block (about 8,000 times for the
+    weights). In a terminal the line is redrawn in place when the whole percentage
+    (or, with no known size, the whole MB) changes. When stderr is not a terminal
+    (CI, log files, notebooks) one line is written per 10%, so logs stay short.
+    """
+
+    def __init__(self) -> None:
+        self.tty = sys.stderr.isatty()
+        self.last: int | None = None
+
+    def __call__(self, blocks: int, block_size: int, total: int) -> None:
+        done = blocks * block_size
+        if total > 0:
+            done = min(done, total)
+            pct = done * 100 // total
+            step = pct if self.tty else pct // 10 * 10
+            text = f"{done / 2**20:.0f} / {total / 2**20:.0f} MB ({step}%)"
+        elif self.tty:
+            step = done // 2**20
+            text = f"{step} MB"
+        else:
+            return
+        if step == self.last:
+            return
+        self.last = step
+        sys.stderr.write(f"\r  {text}" if self.tty else f"  {text}\n")
+        sys.stderr.flush()
+
+    def finish(self) -> None:
+        if self.tty and self.last is not None:
+            sys.stderr.write("\n")
 
 
 def fetch_pretrained() -> Path:
@@ -59,16 +85,17 @@ def fetch_pretrained() -> Path:
     path.parent.mkdir(parents=True, exist_ok=True)
     partial = path.with_name(path.name + ".part")
     print(f"Downloading SubCellSeg weights to {path}", file=sys.stderr)
+    progress = _Progress()
     try:
         try:
-            urllib.request.urlretrieve(url, partial, reporthook=_report)
+            urllib.request.urlretrieve(url, partial, reporthook=progress)
         except OSError as err:  # no network, blocked download, missing file
             raise RuntimeError(
                 f"Could not download the pretrained weights from {url} ({err}).\n"
                 f"Download that file in a web browser and save it as {path}, "
                 "or pass a weights file with --checkpoint."
             ) from None
-        sys.stderr.write("\n")
+        progress.finish()
         if WEIGHTS_SHA256 and not os.environ.get("SUBCELLSEG_WEIGHTS_URL"):
             found = sha256(partial)
             if found != WEIGHTS_SHA256:

tests/test_subcellseg.py

@@ -2,6 +2,7 @@
 they need neither a GPU nor the pretrained weights."""
 
 import csv
+import sys
 
 import numpy as np
 import pytest
@@ -333,3 +334,23 @@ def test_weights_download_is_verified_and_cached(tmp_path, monkeypatch):
     path.unlink()
     with pytest.raises(RuntimeError, match="Could not download"):
         weights.fetch_pretrained()
+
+
+@pytest.mark.parametrize("tty", [False, True])
+def test_download_progress_stays_short(tty, monkeypatch, capsys):
+    from subcellseg import weights
+
+    monkeypatch.setattr(sys.stderr, "isatty", lambda: tty)
+    progress = weights._Progress()
+    total, block = 64 * 2**20, 8192  # the real weights: about 8,000 hook calls
+    for blocks in range(total // block + 2):
+        progress(blocks, block, total)
+    progress.finish()
+    err = capsys.readouterr().err
+    assert "(100%)" in err
+    if tty:
+        assert err.count("\r") == 101  # one redraw per whole percent
+    else:
+        assert "\r" not in err
+        lines = err.splitlines()
+        assert [line.rsplit("(", 1)[1] for line in lines] == [f"{p}%)" for p in range(0, 101, 10)]