[nrf52] Install PlatformIO toolchain Python packages into a dedicated venv (#17635)

esphome/components/nrf52/__init__.py

@@ -69,7 +69,12 @@ from .const import (
     BOOTLOADER_ADAFRUIT_NRF52_SD140_V6,
     BOOTLOADER_ADAFRUIT_NRF52_SD140_V7,
 )
-from .framework import check_and_install, get_build_env, get_build_paths
+from .framework import (
+    check_and_install,
+    get_build_env,
+    get_build_paths,
+    setup_platformio_python_env,
+)
 
 # force import gpio to register pin schema
 from .gpio import nrf52_pin_to_code  # noqa: F401
@@ -514,6 +519,7 @@ def _upload_using_platformio(
 ) -> int | str:
     from esphome.platformio import toolchain
 
+    setup_platformio_python_env()
     if port is not None:
         upload_args += ["--upload-port", port]
     return toolchain.run_platformio_cli_run(config, CORE.verbose, *upload_args)
@@ -809,6 +815,10 @@ def _copy_if_exists(src: Path, dst: Path) -> None:
 
 def run_compile(args, config: ConfigType) -> bool:
     if CORE.using_toolchain_platformio:
+        # The actual build is done by PlatformIO (the caller falls through to
+        # it when this returns False); prepare the Python environment its
+        # Zephyr build script expects first.
+        setup_platformio_python_env()
         return False
     if not CORE.using_toolchain_sdk_nrf:
         raise EsphomeError(

esphome/components/nrf52/framework.py

@@ -4,6 +4,7 @@ import os
 from pathlib import Path
 import platform
 import shutil
+import sys
 import tempfile
 
 import platformdirs
@@ -27,6 +28,11 @@ _LOGGER = logging.getLogger(__name__)
 _REQUIREMENTS = Path(__file__).parent / "requirements.txt"
 TOOLCHAIN_VERSION = "0.17.4"
 
+# Packages the PlatformIO toolchain's Zephyr build script needs beyond west
+# (which comes from requirements.txt). Keep the pin in sync with
+# framework-sdk-nrf scripts/platformio/platformio-build.py.
+_PLATFORMIO_PENV_REQUIREMENTS: tuple[str, ...] = ("cbor2==5.6.5",)
+
 SDK_NG_TOOLCHAIN_MIRRORS = str_to_lst_of_str(
     os.environ.get(
         "ESPHOME_SDK_NG_TOOLCHAIN_MIRRORS",
@@ -145,6 +151,82 @@ def get_build_env() -> dict:
     return env
 
 
+def _get_platformio_penv_path() -> Path:
+    return get_sdk_nrf_tools_path() / "penvs" / "platformio"
+
+
+def _get_penv_site_packages(penv_path: Path) -> Path:
+    if os.name == "nt":
+        return penv_path / "Lib" / "site-packages"
+    python_dir = f"python{sys.version_info.major}.{sys.version_info.minor}"
+    return penv_path / "lib" / python_dir / "site-packages"
+
+
+def _prepend_env_path(name: str, entry: str) -> None:
+    """Prepend ``entry`` to the ``os.pathsep``-separated env var ``name``."""
+    current = os.environ.get(name, "")
+    entries = current.split(os.pathsep) if current else []
+    if entry not in entries:
+        os.environ[name] = os.pathsep.join([entry, *entries])
+
+
+def setup_platformio_python_env() -> None:
+    """Make the Zephyr build's Python packages available to PlatformIO.
+
+    The PlatformIO toolchain's Zephyr framework build script pip-installs
+    west and cbor2 (and pyocd on x86_64) into the Python environment running
+    PlatformIO whenever they are not importable. That environment is not
+    always writable — for example the docker image run as a non-root user,
+    where ESPHome lives in the system Python — so the install fails with
+    "Permission denied". Instead, pre-install those packages into a dedicated
+    venv under the sdk-nrf tools dir and expose it to the PlatformIO
+    subprocesses through the environment:
+
+    * PYTHONPATH makes the venv's packages importable from the interpreter
+      that runs PlatformIO/SCons, so the build script skips its installs.
+    * VIRTUAL_ENV redirects any install the build script still performs via
+      uv (pyocd is fetched on demand) into the writable venv.
+    * PATH exposes console scripts installed into the venv (e.g. pyocd).
+    """
+    penv_path = _get_platformio_penv_path()
+    env_python_path = get_python_env_executable_path(penv_path, "python")
+    sentinel = penv_path / ".ready"
+    # Include the Python version: the venv breaks when the interpreter it
+    # was created from is upgraded, so it must be rebuilt.
+    requirements_hash = hashlib.sha256(
+        _REQUIREMENTS.read_bytes()
+        + "\n".join(_PLATFORMIO_PENV_REQUIREMENTS).encode()
+        + f"python{sys.version_info.major}.{sys.version_info.minor}".encode()
+    ).hexdigest()
+    if (
+        not sentinel.exists()
+        or sentinel.read_text(encoding="utf-8") != requirements_hash
+    ):
+        rmdir(penv_path, msg="Clean up PlatformIO toolchain Python environment")
+
+        create_venv(penv_path, msg="PlatformIO toolchain")
+
+        _LOGGER.info("Installing PlatformIO toolchain requirements ...")
+        cmd = [
+            str(env_python_path),
+            "-m",
+            "pip",
+            "install",
+            "-r",
+            str(_REQUIREMENTS),
+            *_PLATFORMIO_PENV_REQUIREMENTS,
+        ]
+        if not run_command_ok(cmd):
+            raise EsphomeError(
+                "Install requirements for PlatformIO toolchain Python environment failure"
+            )
+        sentinel.write_text(requirements_hash, encoding="utf-8")
+
+    os.environ["VIRTUAL_ENV"] = str(penv_path)
+    _prepend_env_path("PYTHONPATH", str(_get_penv_site_packages(penv_path)))
+    _prepend_env_path("PATH", str(env_python_path.parent))
+
+
 def _patch_uf2conv_escape_sequences(framework_path: Path) -> None:
     # SDK v2.6.1 ships uf2conv.py with '\s+' — an unrecognised escape that
     # Python 3.12+ flags with SyntaxWarning (a future version will reject it).

tests/unit_tests/test_nrf52_framework.py

@@ -3,18 +3,23 @@
 import hashlib
 import os
 from pathlib import Path
+import sys
 from types import SimpleNamespace
 from unittest.mock import patch
 
 import pytest
 
 from esphome.components.nrf52.framework import (
+    _PLATFORMIO_PENV_REQUIREMENTS,
     _REQUIREMENTS,
     TOOLCHAIN_VERSION,
+    _get_penv_site_packages,
+    _get_platformio_penv_path,
     _get_toolchain_platform_info,
     check_and_install,
     get_build_env,
     get_sdk_nrf_tools_path,
+    setup_platformio_python_env,
 )
 from esphome.config_validation import Version
 from esphome.const import KEY_CORE, KEY_FRAMEWORK_VERSION
@@ -255,6 +260,182 @@ class TestCheckAndInstall:
         assert substitutions["extension"] == "tar.xz"
 
 
+# ---------------------------------------------------------------------------
+# setup_platformio_python_env tests
+# ---------------------------------------------------------------------------
+
+
+def _platformio_requirements_hash() -> str:
+    return hashlib.sha256(
+        _REQUIREMENTS.read_bytes()
+        + "\n".join(_PLATFORMIO_PENV_REQUIREMENTS).encode()
+        + f"python{sys.version_info.major}.{sys.version_info.minor}".encode()
+    ).hexdigest()
+
+
+@pytest.fixture
+def platformio_penv_dir() -> Path:
+    """Pre-create the PlatformIO penv dir so sentinel writes succeed.
+
+    create_venv is mocked in these tests, so the directory it would have
+    created must exist for ``sentinel.write_text`` to work.
+    """
+    penv_path = _get_platformio_penv_path()
+    penv_path.mkdir(parents=True, exist_ok=True)
+    return penv_path
+
+
+class TestSetupPlatformioPythonEnv:
+    def test_fresh_install_creates_venv_and_sets_env(
+        self,
+        platformio_penv_dir: Path,
+        mock_nrf52_ops: SimpleNamespace,
+    ) -> None:
+        """No sentinel → venv created, requirements installed, env exported."""
+        with patch.dict(os.environ):
+            os.environ.pop("PYTHONPATH", None)
+
+            setup_platformio_python_env()
+
+            mock_nrf52_ops.rmdir.assert_called_once()
+            mock_nrf52_ops.create_venv.assert_called_once_with(
+                platformio_penv_dir, msg="PlatformIO toolchain"
+            )
+            mock_nrf52_ops.run_command_ok.assert_called_once()
+            cmd = mock_nrf52_ops.run_command_ok.call_args[0][0]
+            assert cmd[1:4] == ["-m", "pip", "install"]
+            assert "-r" in cmd
+            assert str(_REQUIREMENTS) in cmd
+            for requirement in _PLATFORMIO_PENV_REQUIREMENTS:
+                assert requirement in cmd
+            sentinel = platformio_penv_dir / ".ready"
+            assert sentinel.read_text(encoding="utf-8") == (
+                _platformio_requirements_hash()
+            )
+
+            assert os.environ["VIRTUAL_ENV"] == str(platformio_penv_dir)
+            site_packages = str(_get_penv_site_packages(platformio_penv_dir))
+            assert os.environ["PYTHONPATH"] == site_packages
+            bin_dir = str(
+                get_python_env_executable_path(platformio_penv_dir, "python").parent
+            )
+            assert os.environ["PATH"].split(os.pathsep)[0] == bin_dir
+
+    def test_ready_sentinel_skips_install_but_sets_env(
+        self,
+        platformio_penv_dir: Path,
+        mock_nrf52_ops: SimpleNamespace,
+    ) -> None:
+        """Current sentinel → no install work, env vars still exported."""
+        (platformio_penv_dir / ".ready").write_text(
+            _platformio_requirements_hash(), encoding="utf-8"
+        )
+
+        with patch.dict(os.environ):
+            setup_platformio_python_env()
+
+            mock_nrf52_ops.rmdir.assert_not_called()
+            mock_nrf52_ops.create_venv.assert_not_called()
+            mock_nrf52_ops.run_command_ok.assert_not_called()
+            assert os.environ["VIRTUAL_ENV"] == str(platformio_penv_dir)
+
+    def test_stale_sentinel_reinstalls(
+        self,
+        platformio_penv_dir: Path,
+        mock_nrf52_ops: SimpleNamespace,
+    ) -> None:
+        """A sentinel from different requirements → venv rebuilt from scratch."""
+        sentinel = platformio_penv_dir / ".ready"
+        sentinel.write_text("stale-hash", encoding="utf-8")
+
+        with patch.dict(os.environ):
+            setup_platformio_python_env()
+
+        mock_nrf52_ops.rmdir.assert_called_once()
+        mock_nrf52_ops.create_venv.assert_called_once()
+        mock_nrf52_ops.run_command_ok.assert_called_once()
+        assert sentinel.read_text(encoding="utf-8") == _platformio_requirements_hash()
+
+    def test_install_failure_raises(
+        self,
+        platformio_penv_dir: Path,
+        mock_nrf52_ops: SimpleNamespace,
+    ) -> None:
+        """Failing pip install raises EsphomeError and writes no sentinel."""
+        mock_nrf52_ops.run_command_ok.return_value = False
+
+        with (
+            patch.dict(os.environ),
+            pytest.raises(
+                EsphomeError, match="Install requirements for PlatformIO toolchain"
+            ),
+        ):
+            setup_platformio_python_env()
+
+        assert not (platformio_penv_dir / ".ready").exists()
+
+    def test_repeated_calls_do_not_duplicate_env_entries(
+        self,
+        platformio_penv_dir: Path,
+        mock_nrf52_ops: SimpleNamespace,
+    ) -> None:
+        """Compile then upload in one process must not grow PYTHONPATH/PATH."""
+        (platformio_penv_dir / ".ready").write_text(
+            _platformio_requirements_hash(), encoding="utf-8"
+        )
+        site_packages = str(_get_penv_site_packages(platformio_penv_dir))
+        bin_dir = str(
+            get_python_env_executable_path(platformio_penv_dir, "python").parent
+        )
+
+        with patch.dict(os.environ):
+            setup_platformio_python_env()
+            setup_platformio_python_env()
+
+            assert os.environ["PYTHONPATH"].split(os.pathsep).count(site_packages) == 1
+            assert os.environ["PATH"].split(os.pathsep).count(bin_dir) == 1
+
+    def test_existing_pythonpath_preserved(
+        self,
+        platformio_penv_dir: Path,
+        mock_nrf52_ops: SimpleNamespace,
+    ) -> None:
+        """A pre-existing PYTHONPATH keeps its entries after the venv entry."""
+        (platformio_penv_dir / ".ready").write_text(
+            _platformio_requirements_hash(), encoding="utf-8"
+        )
+        site_packages = str(_get_penv_site_packages(platformio_penv_dir))
+
+        with patch.dict(os.environ, {"PYTHONPATH": "/existing/path"}):
+            setup_platformio_python_env()
+
+            assert os.environ["PYTHONPATH"] == os.pathsep.join(
+                [site_packages, "/existing/path"]
+            )
+
+
+@pytest.mark.parametrize(
+    ("os_name", "expected_parts"),
+    [
+        (
+            "posix",
+            (
+                "lib",
+                f"python{sys.version_info.major}.{sys.version_info.minor}",
+                "site-packages",
+            ),
+        ),
+        ("nt", ("Lib", "site-packages")),
+    ],
+)
+def test_get_penv_site_packages(
+    tmp_path: Path, os_name: str, expected_parts: tuple[str, ...]
+) -> None:
+    penv_path = tmp_path / "penv"
+    with patch("os.name", os_name):
+        assert _get_penv_site_packages(penv_path) == penv_path.joinpath(*expected_parts)
+
+
 # ---------------------------------------------------------------------------
 # get_build_env tests
 # ---------------------------------------------------------------------------

tests/unit_tests/test_nrf52_upload.py

@@ -146,6 +146,72 @@ class TestUploadProgramPyocd:
             upload_program(config={}, args=None, host="PYOCD")
 
 
+# ---------------------------------------------------------------------------
+# PlatformIO toolchain paths
+# ---------------------------------------------------------------------------
+
+
+class TestRunCompilePlatformio:
+    def test_prepares_python_env_and_delegates_to_platformio(
+        self, setup_core: Path, tmp_path: Path
+    ) -> None:
+        """The PlatformIO toolchain prepares the env, then returns False so PlatformIO builds."""
+        from esphome.components.nrf52 import run_compile
+
+        _setup_nrf52_core(toolchain=Toolchain.PLATFORMIO, build_path=tmp_path / "build")
+
+        with patch(
+            "esphome.components.nrf52.setup_platformio_python_env"
+        ) as mock_setup:
+            assert run_compile(args=None, config={}) is False
+
+        mock_setup.assert_called_once_with()
+
+
+class TestUploadProgramSerialPlatformio:
+    def _upload(self, host: str, tmp_path: Path, run_result: int) -> tuple:
+        from esphome.components.nrf52 import upload_program
+        from esphome.upload_targets import PortType
+
+        _setup_nrf52_core(toolchain=Toolchain.PLATFORMIO, build_path=tmp_path / "build")
+        CORE.config_path = tmp_path / "test.yaml"
+
+        with (
+            patch("esphome.upload_targets.get_port_type", return_value=PortType.SERIAL),
+            patch("esphome.__main__.check_permissions"),
+            patch("esphome.components.nrf52.setup_platformio_python_env") as mock_setup,
+            patch(
+                "esphome.platformio.toolchain.run_platformio_cli_run",
+                return_value=run_result,
+            ) as mock_run,
+        ):
+            result = upload_program(config={}, args=None, host=host)
+        return result, mock_setup, mock_run
+
+    def test_serial_upload_prepares_env_and_runs_platformio(
+        self, setup_core: Path, tmp_path: Path
+    ) -> None:
+        """Serial upload with the PlatformIO toolchain runs pio with -t upload."""
+        host = "/dev/ttyACM0"
+        result, mock_setup, mock_run = self._upload(host, tmp_path, run_result=0)
+
+        assert result is True
+        mock_setup.assert_called_once_with()
+        mock_run.assert_called_once()
+        run_args = mock_run.call_args[0]
+        assert "-t" in run_args
+        assert "upload" in run_args
+        assert "--upload-port" in run_args
+        assert host in run_args
+
+    def test_serial_upload_failure_raises(
+        self, setup_core: Path, tmp_path: Path
+    ) -> None:
+        """A non-zero PlatformIO result must raise EsphomeError."""
+        with pytest.raises(EsphomeError, match="Upload failed"):
+            self._upload("/dev/ttyACM0", tmp_path, run_result=1)
+
+
 # ---------------------------------------------------------------------------
 # Serial DFU upload path
 # ---------------------------------------------------------------------------