Share tool resolution between the idf_tools helper scripts

This commit is contained in:
J. Nick Koston
2026-08-27 19:21:20 -05:00
parent c1c9408035
commit aaafc1753a
5 changed files with 217 additions and 203 deletions
+68
View File
@@ -0,0 +1,68 @@
"""Shared tool resolution for the sibling idf_tools-backed scripts.
Importable because ``python <script>`` puts this directory first on
sys.path; ``idf_tools`` itself comes from PYTHONPATH.
"""
# pylint: disable=import-error # idf_tools is on PYTHONPATH at runtime only
from collections.abc import Callable, Iterator
import os
from pathlib import Path
from idf_tools import (
CURRENT_PLATFORM,
TOOLS_FILE,
IDFEnv,
ToolBinaryError,
add_and_check_targets,
expand_tools_arg,
g,
load_tools_info,
)
def init_idf_tools(idf_framework_root: str) -> None:
"""Point idf_tools' globals at the framework and IDF_TOOLS_PATH."""
g.idf_path = idf_framework_root
g.idf_tools_path = os.environ.get("IDF_TOOLS_PATH")
g.tools_json = str(Path(g.idf_path) / TOOLS_FILE)
def archive_name(download: object) -> str:
"""The dist/ filename idf_tools downloads and installs this from."""
return download.rename_dist or Path(download.url).name
def iter_tool_downloads(
targets_csv: str,
tool_specs: list[str],
on_broken: Callable[[str, ToolBinaryError], bool],
) -> Iterator[tuple[object, str, str, object]]:
"""Yield (tool, name, version, download) per uninstalled tool, mirroring
``idf_tools.py install``'s expansion; ``on_broken(name, err)`` returns
True to treat a tool with a failing installed binary as not installed."""
targets = add_and_check_targets(IDFEnv.get_idf_env(), targets_csv)
tools_info = load_tools_info()
for name in expand_tools_arg(tool_specs, tools_info, targets):
if "@" in name:
name, version = name.split("@", 1)
else:
version = None
tool = tools_info.get(name)
if tool is None or not tool.compatible_with_platform():
continue
version = version or tool.get_recommended_version()
if version is None:
continue
try:
tool.find_installed_versions()
except ToolBinaryError as e:
if not on_broken(name, e):
continue
if version in tool.versions_installed or version not in tool.versions:
continue
download = tool.versions[version].get_download_for_platform(CURRENT_PLATFORM)
if download is None:
continue
yield tool, name, version, download
+9 -10
View File
@@ -20,6 +20,7 @@ from esphome.build_helpers.pch import ccache_pch_env
from esphome.build_helpers.tools_cache import IDF_TOOLS_CACHE, tools_cache_path
from esphome.core import Version
from esphome.framework_helpers import (
BATCH_EXTRACT_WORKERS,
PathType,
create_venv,
download_and_extract,
@@ -797,21 +798,19 @@ def _preinstall_idf_tool_archives(
tools: list[str],
env: dict[str, str] | None,
) -> None:
"""Extract the prefetched tool archives in parallel before the installer.
``idf_tools.py install`` unpacks one archive at a time on a single core;
``install_tool_archives.py`` drives idf_tools' own ``IDFTool.install()``
with one worker per usable core over the archives the prefetch verified.
The sequential installer still runs afterwards as the authority, skipping
the tools installed here and redoing anything this pass failed on, so
this is strictly best-effort.
"""
"""Run install_tool_archives.py to extract the prefetched tool archives
in parallel. Strictly best-effort: the sequential installer remains the
authority (see that script's docstring)."""
try:
success, _stdout, _stderr = _run_idf_tools_script(
framework_path,
"install_tool_archives.py",
"ESP-IDF tool archive extraction",
args=[targets_str, str(get_usable_cpu_count()), *tools],
args=[
targets_str,
str(min(get_usable_cpu_count(), BATCH_EXTRACT_WORKERS)),
*tools,
],
env=env,
stream_output=True,
)
+22 -53
View File
@@ -18,67 +18,36 @@ or written — this script only reports what the install would download.
from contextlib import redirect_stdout
import json
import os
from pathlib import Path
import sys
from idf_tools import (
CURRENT_PLATFORM,
TOOLS_FILE,
IDFEnv,
ToolBinaryError,
add_and_check_targets,
expand_tools_arg,
g,
get_idf_download_url_apply_mirrors,
load_tools_info,
)
from _tool_resolution import archive_name, init_idf_tools, iter_tool_downloads
from idf_tools import ToolBinaryError, get_idf_download_url_apply_mirrors
def collect_downloads() -> list[dict]:
g.idf_path = sys.argv[1]
g.idf_tools_path = os.environ.get("IDF_TOOLS_PATH")
g.tools_json = str(Path(g.idf_path) / TOOLS_FILE)
init_idf_tools(sys.argv[1])
targets = add_and_check_targets(IDFEnv.get_idf_env(), sys.argv[2])
tools_info = load_tools_info()
downloads: list[dict] = []
def on_broken(name: str, e: ToolBinaryError) -> bool:
# A broken installed binary is idf_tools' problem to repair on
# install; note it and treat the version as not installed.
print(f"tool {name} failed its binary check: {e}", file=sys.stderr)
return True
for name in expand_tools_arg(sys.argv[3:], tools_info, targets):
if "@" in name:
name, version = name.split("@", 1)
else:
version = None
tool = tools_info.get(name)
if tool is None or not tool.compatible_with_platform():
continue
version = version or tool.get_recommended_version()
if version is None:
continue
try:
tool.find_installed_versions()
except ToolBinaryError as e:
# A broken installed binary is idf_tools' problem to repair on
# install; note it and treat the version as not installed.
print(f"tool {name} failed its binary check: {e}", file=sys.stderr)
if version in tool.versions_installed or version not in tool.versions:
continue
download = tool.versions[version].get_download_for_platform(CURRENT_PLATFORM)
if download is None:
continue
downloads.append(
{
"name": f"{name}@{version}",
# Apply the same IDF_MIRROR_PREFIX_MAP / IDF_GITHUB_ASSETS
# rewriting the installer's own downloader applies, so users
# behind a mirror prefetch from the mirror too.
"url": get_idf_download_url_apply_mirrors(None, download.url),
"size": download.size,
"sha256": download.sha256,
"dest": download.rename_dist or Path(download.url).name,
}
return [
{
"name": f"{name}@{version}",
# Apply the same IDF_MIRROR_PREFIX_MAP / IDF_GITHUB_ASSETS
# rewriting the installer's own downloader applies, so users
# behind a mirror prefetch from the mirror too.
"url": get_idf_download_url_apply_mirrors(None, download.url),
"size": download.size,
"sha256": download.sha256,
"dest": archive_name(download),
}
for _tool, name, version, download in iter_tool_downloads(
sys.argv[2], sys.argv[3:], on_broken
)
return downloads
]
# idf_tools prints informational lines (e.g. mirror URL rewrites) to stdout;
+31 -63
View File
@@ -1,82 +1,47 @@
"""Extract prefetched ESP-IDF tool archives in parallel.
Run via ``python <this file> <idf_framework_root> <targets-csv> <workers>
<tool-spec>...``. PYTHONPATH must include ``<idf_framework_root>/tools`` so
``idf_tools`` is importable, and IDF_TOOLS_PATH must be set.
``idf_tools.py install`` unpacks one archive at a time; this extracts every
tool whose verified archive the prefetch already placed in
``<IDF_TOOLS_PATH>/dist``, several at once, using idf_tools' own
``IDFTool.install()`` so unpacking, container-dir stripping, and the binary
check match the sequential installer exactly. That installer still runs
afterwards as the authority: it skips the tools installed here and redoes
anything this pass failed on, so per-tool failures only warn on stderr.
<tool-spec>...`` with idf_tools on PYTHONPATH and IDF_TOOLS_PATH set.
Drives idf_tools' own ``IDFTool.install()`` so extraction semantics match
the sequential installer, which still runs afterwards as the authority and
redoes anything this best-effort pass failed on.
"""
# pylint: disable=import-error # idf_tools is on PYTHONPATH at runtime only
from concurrent.futures import ThreadPoolExecutor
import os
from pathlib import Path
import sys
from idf_tools import (
CURRENT_PLATFORM,
TOOLS_FILE,
IDFEnv,
ToolBinaryError,
add_and_check_targets,
expand_tools_arg,
g,
load_tools_info,
)
from _tool_resolution import archive_name, init_idf_tools, iter_tool_downloads
from idf_tools import ToolBinaryError, g
def collect_pending() -> list[tuple[object, str, str]]:
"""The (tool, name, version) jobs whose verified archive is on disk."""
g.idf_path = sys.argv[1]
g.idf_tools_path = os.environ.get("IDF_TOOLS_PATH")
g.tools_json = str(Path(g.idf_path) / TOOLS_FILE)
targets = add_and_check_targets(IDFEnv.get_idf_env(), sys.argv[2])
tools_info = load_tools_info()
def collect_pending(
targets_csv: str, tool_specs: list[str]
) -> dict[tuple[str, str], object]:
"""The {(name, version): tool} jobs whose verified archive is on disk."""
dist_path = Path(g.idf_tools_path) / "dist"
pending: list[tuple[object, str, str]] = []
seen: set[tuple[str, str]] = set()
for name in expand_tools_arg(sys.argv[4:], tools_info, targets):
if "@" in name:
name, version = name.split("@", 1)
else:
version = None
tool = tools_info.get(name)
if tool is None or not tool.compatible_with_platform():
continue
version = version or tool.get_recommended_version()
if version is None:
continue
try:
tool.find_installed_versions()
except ToolBinaryError as e:
# Repairing a broken installed binary is the installer's job
print(f"leaving broken {name} to the installer: {e}", file=sys.stderr)
continue
if version in tool.versions_installed or version not in tool.versions:
continue
download = tool.versions[version].get_download_for_platform(CURRENT_PLATFORM)
if download is None:
continue
def on_broken(name: str, e: ToolBinaryError) -> bool:
# Repairing a broken installed binary is the installer's job
print(f"leaving broken {name} to the installer: {e}", file=sys.stderr)
return False
pending: dict[tuple[str, str], object] = {}
for tool, name, version, download in iter_tool_downloads(
targets_csv, tool_specs, on_broken
):
# An archive at its final name was sha256-verified by the prefetch
archive = dist_path / (download.rename_dist or Path(download.url).name)
if not archive.is_file() or (name, version) in seen:
if (name, version) in pending or not (
dist_path / archive_name(download)
).is_file():
continue
seen.add((name, version))
pending.append((tool, name, version))
pending[(name, version)] = tool
return pending
def install_one(job: tuple[object, str, str]) -> None:
tool, name, version = job
def install_one(tool: object, name: str, version: str) -> None:
try:
tool.install(version)
# check_binary_valid exits via SystemExit; the installer redoes failures
@@ -88,19 +53,22 @@ def install_one(job: tuple[object, str, str]) -> None:
def main() -> None:
pending = collect_pending()
_script, idf_framework_root, targets_csv, workers_str, *tool_specs = sys.argv
init_idf_tools(idf_framework_root)
pending = collect_pending(targets_csv, tool_specs)
if len(pending) < 2:
# Nothing to parallelize; the installer keeps its normal output
return
workers = min(int(sys.argv[3]), len(pending))
workers = min(int(workers_str), len(pending))
print(
f"Extracting {len(pending)} ESP-IDF tool archive(s) with "
f"{workers} worker(s): "
+ ", ".join(f"{name}@{version}" for _, name, version in pending),
+ ", ".join(f"{name}@{version}" for name, version in pending),
flush=True,
)
with ThreadPoolExecutor(max_workers=workers) as ex:
list(ex.map(install_one, pending))
for (name, version), tool in pending.items():
ex.submit(install_one, tool, name, version)
main()
+87 -77
View File
@@ -1199,13 +1199,16 @@ def test_framework_install_prefetches_before_installer(
_IDF_TOOLS_STUB_DIR = Path(__file__).parent / "fixtures" / "idf_tools_stub"
_ESPIDF_SCRIPTS_DIR = Path(__file__).parents[2] / "esphome" / "espidf"
def _run_downloads_script(
tmp_path: Path, *args: str, env_extra: dict[str, str] | None = None
def _run_espidf_script(
tmp_path: Path,
script_name: str,
*args: str,
env_extra: dict[str, str] | None = None,
) -> subprocess.CompletedProcess[str]:
"""Run the real get_tool_downloads.py against the stub idf_tools module."""
script = Path(__file__).parents[2] / "esphome" / "espidf" / "get_tool_downloads.py"
"""Run a real espidf helper script against the stub idf_tools module."""
env = os.environ | {
"PYTHONPATH": str(_IDF_TOOLS_STUB_DIR),
"IDF_TOOLS_PATH": str(tmp_path / "tp"),
@@ -1213,7 +1216,12 @@ def _run_downloads_script(
if env_extra:
env |= env_extra
return subprocess.run(
[sys.executable, str(script), str(tmp_path / "fw"), *args],
[
sys.executable,
str(_ESPIDF_SCRIPTS_DIR / script_name),
str(tmp_path / "fw"),
*args,
],
capture_output=True,
text=True,
env=env,
@@ -1225,7 +1233,7 @@ def test_get_tool_downloads_lists_missing_tools(tmp_path: Path) -> None:
"""Installed versions are skipped, tools that fail their binary check are
still listed, rename_dist decides the dist filename, and idf_tools' stdout
chatter stays off the JSON channel."""
result = _run_downloads_script(tmp_path, "esp32", "required")
result = _run_espidf_script(tmp_path, "get_tool_downloads.py", "esp32", "required")
assert result.returncode == 0, result.stderr
downloads = {d["name"]: d for d in json.loads(result.stdout)}
@@ -1241,8 +1249,9 @@ def test_get_tool_downloads_lists_missing_tools(tmp_path: Path) -> None:
def test_get_tool_downloads_applies_mirror_rewrite(tmp_path: Path) -> None:
result = _run_downloads_script(
result = _run_espidf_script(
tmp_path,
"get_tool_downloads.py",
"esp32",
"required",
env_extra={"TEST_MIRROR_PREFIX": "https://mirror.test/"},
@@ -1253,13 +1262,13 @@ def test_get_tool_downloads_applies_mirror_rewrite(tmp_path: Path) -> None:
assert all(d["url"].startswith("https://mirror.test/") for d in downloads)
def _run_downloads_inprocess(
def _run_espidf_script_inprocess(
tmp_path: Path,
monkeypatch: pytest.MonkeyPatch,
capsys: pytest.CaptureFixture[str],
script_name: str,
*args: str,
) -> list[dict]:
"""Execute get_tool_downloads.py in-process against the stub idf_tools.
) -> None:
"""Execute an espidf helper script in-process against the stub idf_tools.
Unlike the subprocess variant this runs under coverage, exercising the
script's own lines.
@@ -1270,10 +1279,24 @@ def _run_downloads_inprocess(
stub = importlib.util.module_from_spec(spec)
spec.loader.exec_module(stub)
monkeypatch.setitem(sys.modules, "idf_tools", stub)
# _tool_resolution binds idf_tools objects at import; force a fresh
# import against this test's stub instance
monkeypatch.delitem(sys.modules, "_tool_resolution", raising=False)
# python <script> puts the script's directory on sys.path; runpy does not
monkeypatch.syspath_prepend(str(_ESPIDF_SCRIPTS_DIR))
monkeypatch.setenv("IDF_TOOLS_PATH", str(tmp_path / "tp"))
script = Path(__file__).parents[2] / "esphome" / "espidf" / "get_tool_downloads.py"
script = _ESPIDF_SCRIPTS_DIR / script_name
monkeypatch.setattr(sys, "argv", [str(script), str(tmp_path / "fw"), *args])
runpy.run_path(str(script))
def _run_downloads_inprocess(
tmp_path: Path,
monkeypatch: pytest.MonkeyPatch,
capsys: pytest.CaptureFixture[str],
*args: str,
) -> list[dict]:
_run_espidf_script_inprocess(tmp_path, monkeypatch, "get_tool_downloads.py", *args)
return json.loads(capsys.readouterr().out)
@@ -2017,6 +2040,19 @@ def test_preinstall_streams_script_with_workers(tmp_path: Path) -> None:
)
def test_preinstall_caps_workers(tmp_path: Path) -> None:
"""A high core count is capped; the workers share one disk."""
with (
patch(
"esphome.espidf.framework._run_idf_tools_script",
return_value=(True, None, None),
) as run_script,
patch("esphome.espidf.framework.get_usable_cpu_count", return_value=64),
):
_preinstall_idf_tool_archives(tmp_path, "esp32", ["required"], None)
assert run_script.call_args.kwargs["args"][1] == "10"
def test_preinstall_script_failure_only_warns(
tmp_path: Path, caplog: pytest.LogCaptureFixture
) -> None:
@@ -2052,28 +2088,6 @@ def test_preinstall_exception_only_warns(
# ---------------------------------------------------------------------------
def _run_install_script(
tmp_path: Path, *args: str, env_extra: dict[str, str] | None = None
) -> subprocess.CompletedProcess[str]:
"""Run the real install_tool_archives.py against the stub idf_tools."""
script = (
Path(__file__).parents[2] / "esphome" / "espidf" / "install_tool_archives.py"
)
env = os.environ | {
"PYTHONPATH": str(_IDF_TOOLS_STUB_DIR),
"IDF_TOOLS_PATH": str(tmp_path / "tp"),
}
if env_extra:
env |= env_extra
return subprocess.run(
[sys.executable, str(script), str(tmp_path / "fw"), *args],
capture_output=True,
text=True,
env=env,
check=False,
)
def _make_dist(tmp_path: Path, *names: str) -> None:
dist = tmp_path / "tp" / "dist"
dist.mkdir(parents=True, exist_ok=True)
@@ -2082,10 +2096,12 @@ def _make_dist(tmp_path: Path, *names: str) -> None:
def test_install_tool_archives_extracts_pending_in_parallel(tmp_path: Path) -> None:
"""Tools with a prefetched archive install concurrently; installed tools,
broken tools, and tools without an archive stay with the installer."""
"""Subprocess end-to-end: tools with a prefetched archive install
concurrently; installed and broken tools stay with the installer."""
_make_dist(tmp_path, "cmake.tar.gz", "ninja-v1.zip", "x.tar.gz", "y.tar.gz")
result = _run_install_script(tmp_path, "esp32", "8", "required")
result = _run_espidf_script(
tmp_path, "install_tool_archives.py", "esp32", "8", "required"
)
assert result.returncode == 0, result.stderr
tools = tmp_path / "tp" / "tools"
assert (tools / "cmake" / "3.30.2" / ".installed").is_file()
@@ -2102,71 +2118,65 @@ def test_install_tool_archives_extracts_pending_in_parallel(tmp_path: Path) -> N
def test_install_tool_archives_single_pending_stays_sequential(
tmp_path: Path,
monkeypatch: pytest.MonkeyPatch,
capsys: pytest.CaptureFixture[str],
) -> None:
"""One pending archive has nothing to parallelize; the installer keeps
its normal output."""
_make_dist(tmp_path, "cmake.tar.gz")
result = _run_install_script(tmp_path, "esp32", "4", "required")
assert result.returncode == 0, result.stderr
_run_espidf_script_inprocess(
tmp_path, monkeypatch, "install_tool_archives.py", "esp32", "4", "required"
)
assert not (tmp_path / "tp" / "tools").exists()
assert "Extracting" not in result.stdout
assert "Extracting" not in capsys.readouterr().out
def test_install_tool_archives_failed_install_left_to_installer(
tmp_path: Path,
monkeypatch: pytest.MonkeyPatch,
capsys: pytest.CaptureFixture[str],
) -> None:
"""A per-tool failure (SystemExit from the binary check) warns and moves
on; the other tools still install."""
_make_dist(tmp_path, "cmake.tar.gz", "ninja-v1.zip")
result = _run_install_script(
tmp_path, "esp32", "4", "required", env_extra={"TEST_FAIL_INSTALL": "ninja"}
monkeypatch.setenv("TEST_FAIL_INSTALL", "ninja")
_run_espidf_script_inprocess(
tmp_path, monkeypatch, "install_tool_archives.py", "esp32", "4", "required"
)
assert result.returncode == 0, result.stderr
tools = tmp_path / "tp" / "tools"
assert (tools / "cmake" / "3.30.2" / ".installed").is_file()
assert not (tools / "ninja").exists()
assert "pre-extracting ninja@1.12.1 failed" in result.stderr
assert "pre-extracting ninja@1.12.1 failed" in capsys.readouterr().err
def _run_install_inprocess(
tmp_path: Path,
monkeypatch: pytest.MonkeyPatch,
*args: str,
) -> None:
"""Execute install_tool_archives.py in-process against the stub idf_tools.
Unlike the subprocess variant this runs under coverage, exercising the
script's own lines.
"""
spec = importlib.util.spec_from_file_location(
"idf_tools", _IDF_TOOLS_STUB_DIR / "idf_tools.py"
)
stub = importlib.util.module_from_spec(spec)
spec.loader.exec_module(stub)
monkeypatch.setitem(sys.modules, "idf_tools", stub)
monkeypatch.setenv("IDF_TOOLS_PATH", str(tmp_path / "tp"))
script = (
Path(__file__).parents[2] / "esphome" / "espidf" / "install_tool_archives.py"
)
monkeypatch.setattr(sys, "argv", [str(script), str(tmp_path / "fw"), *args])
runpy.run_path(str(script))
def test_install_tool_archives_inprocess_dedupes_specs(
def test_install_tool_archives_inprocess_dedupes_and_skips(
tmp_path: Path,
monkeypatch: pytest.MonkeyPatch,
capsys: pytest.CaptureFixture[str],
) -> None:
"""In-process full flow: duplicate tool@version specs collapse to one
job and both tools install."""
_make_dist(tmp_path, "cmake.tar.gz", "ninja-v1.zip")
_run_install_inprocess(
tmp_path, monkeypatch, "esp32", "8", "cmake", "ninja", "cmake@3.30.2"
job, broken and installed tools are skipped, both pending tools install."""
_make_dist(tmp_path, "cmake.tar.gz", "ninja-v1.zip", "x.tar.gz", "y.tar.gz")
_run_espidf_script_inprocess(
tmp_path,
monkeypatch,
"install_tool_archives.py",
"esp32",
"8",
"cmake",
"ninja",
"cmake@3.30.2",
"installed-tool",
"broken-tool",
)
out = capsys.readouterr().out
captured = capsys.readouterr()
assert (
"Extracting 2 ESP-IDF tool archive(s) with 2 worker(s): "
"cmake@3.30.2, ninja@1.12.1" in out
"cmake@3.30.2, ninja@1.12.1" in captured.out
)
assert (tmp_path / "tp" / "tools" / "cmake" / "3.30.2" / ".installed").is_file()
assert (tmp_path / "tp" / "tools" / "ninja" / "1.12.1" / ".installed").is_file()
assert "leaving broken broken-tool to the installer" in captured.err
tools = tmp_path / "tp" / "tools"
assert (tools / "cmake" / "3.30.2" / ".installed").is_file()
assert (tools / "ninja" / "1.12.1" / ".installed").is_file()
assert not (tools / "installed-tool").exists()
assert not (tools / "broken-tool").exists()