Merge branch 'esp8266-native-build-infra' into esp8266-native-framework-installer

esphome/platformio/registry.py

@@ -4,6 +4,7 @@ platformio package (identical bits, esphome's own download machinery)."""
 from __future__ import annotations
 
 from collections.abc import Collection
+from concurrent.futures import ThreadPoolExecutor
 import io
 import json
 import logging
@@ -13,6 +14,7 @@ import platform
 
 from esphome.core import EsphomeError
 from esphome.framework_helpers import (
+    BatchDownloadProgress,
     archive_extract_all,
     download_from_mirrors,
     download_with_resume,
@@ -140,6 +142,68 @@ def _check_layout(name: str, dest: Path, expect: Collection[str]) -> None:
             )
 
 
+def prefetch_packages(
+    packages: list[tuple[str, str, Path, list[str]]], downloads_dir: Path
+) -> None:
+    """Download pending package archives in parallel under one combined bar.
+
+    ``packages`` holds ``(name, version, dest, mirrors)`` per package. Purely
+    an optimization: ``install_package`` verifies every archive and
+    re-downloads anything this pass left unfinished. Mirror overrides and
+    registry entries without a size stay on the sequential path so its
+    per-file bars remain trustworthy.
+    """
+    pending: list[tuple[str, str, str, str, int]] = []
+    for name, version, dest, mirrors in packages:
+        if mirrors or (dest / ".esphome_extracted").is_file():
+            continue
+        try:
+            url, sha256, size = registry_download(name, version)
+        except EsphomeError as err:
+            # The sequential install reports the real failure with context
+            _LOGGER.debug("Prefetch resolve for %s failed: %s", name, err)
+            continue
+        if not size:
+            continue
+        archive = downloads_dir / f"{name}-{version}"
+        if archive.is_file() and archive.stat().st_size == size:
+            continue
+        pending.append((name, version, url, sha256, size))
+    if len(pending) < 2:
+        return
+    downloads_dir.mkdir(parents=True, exist_ok=True)
+    _LOGGER.info(
+        "Downloading %d package archive(s): %s",
+        len(pending),
+        ", ".join(name for name, _, _, _, _ in pending),
+    )
+    progress = BatchDownloadProgress(
+        "Downloading packages", sum(size for *_, size in pending)
+    )
+
+    def _fetch(entry: tuple[str, str, str, str, int]) -> None:
+        name, version, url, sha256, size = entry
+        try:
+            download_with_resume(
+                url,
+                downloads_dir / f"{name}-{version}",
+                sha256=sha256,
+                size=size,
+                progress=progress.tracker(),
+            )
+        except Exception as err:  # noqa: BLE001  # pylint: disable=broad-exception-caught
+            # install_package retries this one itself, with a visible bar
+            _LOGGER.debug("Prefetch of %s failed: %s", name, err)
+
+    ex = ThreadPoolExecutor(max_workers=len(pending))
+    try:
+        for future in [ex.submit(_fetch, entry) for entry in pending]:
+            future.result()
+    finally:
+        ex.shutdown(wait=True, cancel_futures=True)
+        progress.done()
+
+
 def install_package(
     name: str,
     version: str,

tests/unit_tests/test_platformio_registry.py

@@ -452,3 +452,155 @@ def test_registry_download_non_list_system_is_named() -> None:
         pytest.raises(EsphomeError, match="Unexpected package registry response"),
     ):
         registry.registry_download("pkg", "1.0.0")
+
+
+def _resolve_for(sizes: dict[str, int | None]):
+    def resolve(name: str, version: str):
+        size = sizes[name]
+        if size == -1:
+            raise EsphomeError("registry down")
+        return (f"http://x/{name}.tar.gz", "abc123", size)
+
+    return resolve
+
+
+def test_prefetch_packages_downloads_pending_in_parallel(tmp_path: Path) -> None:
+    """Two uninstalled packages download together under one combined bar,
+    with the registry's sha256 and size and a batch progress tracker."""
+    with (
+        patch.object(registry, "download_with_resume") as mock_download,
+        patch.object(
+            registry, "registry_download", side_effect=_resolve_for({"a": 10, "b": 20})
+        ),
+    ):
+        registry.prefetch_packages(
+            [
+                ("a", "1.0", tmp_path / "a", []),
+                ("b", "2.0", tmp_path / "b", []),
+            ],
+            tmp_path / "dl",
+        )
+    assert mock_download.call_count == 2
+    for call, (name, version, size) in zip(
+        mock_download.call_args_list, [("a", "1.0", 10), ("b", "2.0", 20)], strict=True
+    ):
+        assert call[0][0] == f"http://x/{name}.tar.gz"
+        assert call[0][1] == tmp_path / "dl" / f"{name}-{version}"
+        assert call[1]["sha256"] == "abc123"
+        assert call[1]["size"] == size
+        assert callable(call[1]["progress"])
+
+
+def test_prefetch_packages_single_pending_skips(tmp_path: Path) -> None:
+    """One pending package has nothing to parallelize; the sequential
+    install keeps its own bar."""
+    marker_dest = tmp_path / "a"
+    marker_dest.mkdir()
+    (marker_dest / ".esphome_extracted").touch()
+    with (
+        patch.object(registry, "download_with_resume") as mock_download,
+        patch.object(
+            registry, "registry_download", side_effect=_resolve_for({"b": 20})
+        ),
+    ):
+        registry.prefetch_packages(
+            [
+                ("a", "1.0", marker_dest, []),
+                ("b", "2.0", tmp_path / "b", []),
+            ],
+            tmp_path / "dl",
+        )
+    mock_download.assert_not_called()
+
+
+def test_prefetch_packages_mirror_and_sizeless_stay_sequential(
+    tmp_path: Path,
+) -> None:
+    """Mirror overrides and size-less registry entries are left to the
+    sequential path so its per-file bars stay trustworthy."""
+    with (
+        patch.object(registry, "download_with_resume") as mock_download,
+        patch.object(
+            registry,
+            "registry_download",
+            side_effect=_resolve_for({"b": None, "c": 30}),
+        ),
+    ):
+        registry.prefetch_packages(
+            [
+                ("a", "1.0", tmp_path / "a", ["http://mirror/{VERSION}"]),
+                ("b", "2.0", tmp_path / "b", []),
+                ("c", "3.0", tmp_path / "c", []),
+            ],
+            tmp_path / "dl",
+        )
+    mock_download.assert_not_called()
+
+
+def test_prefetch_packages_resolve_failure_defers_to_install(
+    tmp_path: Path, caplog: pytest.LogCaptureFixture
+) -> None:
+    """A registry failure only skips the prefetch; install_package reports
+    the real error with context."""
+    caplog.set_level("DEBUG")
+    with (
+        patch.object(registry, "download_with_resume") as mock_download,
+        patch.object(
+            registry, "registry_download", side_effect=_resolve_for({"a": -1, "b": 20})
+        ),
+    ):
+        registry.prefetch_packages(
+            [
+                ("a", "1.0", tmp_path / "a", []),
+                ("b", "2.0", tmp_path / "b", []),
+            ],
+            tmp_path / "dl",
+        )
+    mock_download.assert_not_called()
+    assert "Prefetch resolve for a failed" in caplog.text
+
+
+def test_prefetch_packages_complete_archive_skipped(tmp_path: Path) -> None:
+    """An archive already fully downloaded is not re-fetched."""
+    dl = tmp_path / "dl"
+    dl.mkdir()
+    (dl / "a-1.0").write_bytes(b"x" * 10)
+    with (
+        patch.object(registry, "download_with_resume") as mock_download,
+        patch.object(
+            registry, "registry_download", side_effect=_resolve_for({"a": 10, "b": 20})
+        ),
+    ):
+        registry.prefetch_packages(
+            [
+                ("a", "1.0", tmp_path / "a", []),
+                ("b", "2.0", tmp_path / "b", []),
+            ],
+            dl,
+        )
+    mock_download.assert_not_called()
+
+
+def test_prefetch_packages_download_failure_is_debug(
+    tmp_path: Path, caplog: pytest.LogCaptureFixture
+) -> None:
+    """A failed prefetch download is logged and left for install_package."""
+    caplog.set_level("DEBUG")
+    with (
+        patch.object(
+            registry, "download_with_resume", side_effect=OSError("boom")
+        ) as mock_download,
+        patch.object(
+            registry, "registry_download", side_effect=_resolve_for({"a": 10, "b": 20})
+        ),
+    ):
+        registry.prefetch_packages(
+            [
+                ("a", "1.0", tmp_path / "a", []),
+                ("b", "2.0", tmp_path / "b", []),
+            ],
+            tmp_path / "dl",
+        )
+    assert mock_download.call_count == 2
+    assert "Prefetch of a failed" in caplog.text
+    assert "Prefetch of b failed" in caplog.text