Finish the batch bar, cancel queued archives on Ctrl-C, report known sizes

esphome/espidf/framework.py

@@ -747,27 +747,38 @@ def _prefetch_idf_tool_archives(
             len(entries),
             ", ".join(entry["name"] for entry in entries),
         )
+        # tools.json always carries sizes; should one be missing the bar
+        # could not be trusted, so draw none rather than a wrong one.
+        sizes = [entry["size"] for entry in entries]
         progress = BatchDownloadProgress(
-            "Downloading ESP-IDF tools",
-            sum(entry["size"] or 0 for entry in entries),
+            "Downloading ESP-IDF tools", sum(sizes) if all(sizes) else 0
         )
 
         def _download(entry: dict) -> None:
+            tracker = progress.tracker()
             try:
                 download_with_resume(
                     entry["url"],
                     dist_path / entry["dest"],
                     sha256=entry["sha256"],
                     size=entry["size"],
-                    progress=progress.tracker(),
+                    progress=tracker,
                 )
             except Exception as e:  # noqa: BLE001  # pylint: disable=broad-exception-caught
                 # Keep prefetching the remaining archives; the installer
                 # will retry this one itself (without resume).
+                tracker(0)
                 _LOGGER.warning("Could not prefetch %s: %s", entry["name"], e)
 
-        with ThreadPoolExecutor(max_workers=min(_PREFETCH_WORKERS, len(entries))) as ex:
-            list(ex.map(_download, entries))
+        ex = ThreadPoolExecutor(max_workers=min(_PREFETCH_WORKERS, len(entries)))
+        try:
+            for future in [ex.submit(_download, entry) for entry in entries]:
+                future.result()
+        finally:
+            # On Ctrl-C drop the queued archives instead of downloading them
+            # all before the process can exit; in-flight ones still finish.
+            ex.shutdown(wait=True, cancel_futures=True)
+            progress.done()
     except Exception as e:  # noqa: BLE001  # pylint: disable=broad-exception-caught
         # The installer downloads anything missing itself; never let the
         # prefetch become a new way for the install to fail.

esphome/framework_helpers.py

@@ -742,7 +742,9 @@ class BatchDownloadProgress:
     Each ``tracker()`` is a ``progress`` callback for one download; it reports
     that file's absolute byte count and the bar shows the sum over ``total``.
     The lock also serialises the bar's stderr writes, so worker threads never
-    interleave frames. With an unknown ``total`` (0) nothing is drawn.
+    interleave frames. With an unknown ``total`` (0) nothing is drawn. Call
+    ``done()`` once every download has finished (or failed) so a bar that
+    never reached 100% still ends its line before the next log message.
     """
 
     def __init__(self, header: str, total: int) -> None:
@@ -765,6 +767,10 @@ class BatchDownloadProgress:
 
         return update
 
+    def done(self) -> None:
+        if self._bar is not None and self._bar.last_progress != 100:
+            self._bar.done()
+
 
 def download_with_resume(
     url: str,
@@ -829,7 +835,7 @@ def download_with_resume(
         try:
             _verify_file(dest, sha256, size)
             if progress is not None:
-                progress(dest.stat().st_size)
+                progress(size if size is not None else dest.stat().st_size)
             return
         except EsphomeError:
             dest.unlink()
@@ -887,7 +893,7 @@ def download_with_resume(
             if progress is not None:
                 # Also credits a part file an earlier run completed without
                 # streaming anything this time.
-                progress(part.stat().st_size)
+                progress(expected_size or part.stat().st_size)
             if not expected_size and sha256 is None:
                 # No sha, no size, and the server sent no usable
                 # content-length: nothing can prove the download complete

tests/unit_tests/test_espidf_framework.py

@@ -15,7 +15,7 @@ import subprocess
 import sys
 import tarfile
 from types import SimpleNamespace
-from unittest.mock import patch
+from unittest.mock import MagicMock, patch
 
 import pytest
 
@@ -1043,6 +1043,29 @@ def test_prefetch_one_failed_archive_does_not_stop_the_rest(
     assert "Could not prefetch cmake@3.30.2" in caplog.text
 
 
+def test_prefetch_finishes_progress_bar_and_cancels_queue(tmp_path: Path) -> None:
+    """The batch bar is closed out after the pool, and the pool is shut down
+    with cancel_futures so Ctrl-C does not drain every queued archive."""
+    with (
+        patch(
+            "esphome.espidf.framework.run_command",
+            return_value=(True, _PREFETCH_JSON, ""),
+        ),
+        patch("esphome.espidf.framework.download_with_resume"),
+        patch("esphome.espidf.framework.get_system_python_path", return_value="python"),
+        patch("esphome.espidf.framework.BatchDownloadProgress") as progress_cls,
+        patch(
+            "esphome.espidf.framework.ThreadPoolExecutor", wraps=ThreadPoolExecutor
+        ) as pool_cls,
+    ):
+        pool = MagicMock(wraps=ThreadPoolExecutor(max_workers=2))
+        pool_cls.return_value = pool
+        _prefetch_idf_tool_archives(tmp_path, "esp32", ["required"], None)
+
+    pool.shutdown.assert_called_once_with(wait=True, cancel_futures=True)
+    progress_cls.return_value.done.assert_called_once_with()
+
+
 def test_prefetch_passes_targets_and_tools_to_script(tmp_path: Path) -> None:
     with (
         patch(

tests/unit_tests/test_framework_helpers.py

@@ -1181,9 +1181,31 @@ class TestBatchDownloadProgress:
 
     def test_unknown_total_draws_nothing(self) -> None:
         with patch("esphome.framework_helpers.ProgressBar") as bar_cls:
-            BatchDownloadProgress("Downloading", 0).tracker()(5)
+            progress = BatchDownloadProgress("Downloading", 0)
+            progress.tracker()(5)
+            progress.done()
         bar_cls.assert_not_called()
 
+    def test_done_ends_an_unfinished_bar(self) -> None:
+        """A batch that stops short of 100% (a failed archive) still ends its
+        line so the next log message starts on a fresh row."""
+        stream = io.StringIO()
+        stream.isatty = lambda: True  # type: ignore[method-assign]
+        with patch("esphome.helpers.sys.stderr", stream):
+            progress = BatchDownloadProgress("Downloading", 10)
+            progress.tracker()(5)
+            progress.done()
+        assert stream.getvalue().endswith("50% \n")
+
+    def test_done_after_full_bar_adds_nothing(self) -> None:
+        stream = io.StringIO()
+        stream.isatty = lambda: True  # type: ignore[method-assign]
+        with patch("esphome.helpers.sys.stderr", stream):
+            progress = BatchDownloadProgress("Downloading", 10)
+            progress.tracker()(10)
+            progress.done()
+        assert stream.getvalue().endswith("100% Done...\r\n")
+
 
 class TestDownloadFromMirrors:
     def test_success_returns_url_and_writes_content(self, tmp_path: Path) -> None: