zephyr: zmake: add re-exec logic

Add some logic to re-exec zmake when running inside of a chroot and we
find the source code available (to reduce the number of times we have
to update the chroot).

Note: we also discussed adding a chroot-upgrade-required warning
logic.  This is actually a separate idea so will be handled in a
separate CL.

BUG=b:197636145
BRANCH=none
TEST=unit tests

Signed-off-by: Jack Rosenthal <jrosenth@chromium.org>
Change-Id: I26ecbb9a575f22c8667a1928e4bd5836f6fd4fe1
Reviewed-on: https://chromium-review.googlesource.com/c/chromiumos/platform/ec/+/3114503
Reviewed-by: Yuval Peress <peress@chromium.org>

2 files changed, 111 insertions, 0 deletions

diff --git a/zephyr/zmake/tests/test_reexec.py b/zephyr/zmake/tests/test_reexec.py
new file mode 100644
index 0000000000..9f25b5a834
--- /dev/null
+++ b/zephyr/zmake/tests/test_reexec.py
@@ -0,0 +1,59 @@
+# Copyright 2021 The Chromium OS Authors. All rights reserved.
+# Use of this source code is governed by a BSD-style license that can be
+# found in the LICENSE file.
+"""Test the zmake re-exec functionality."""
+
+import os
+import sys
+import unittest.mock as mock
+
+import pytest
+
+import zmake.__main__ as main
+
+
+@pytest.fixture
+def fake_env(monkeypatch):
+    environ = {}
+    monkeypatch.setattr(os, "environ", environ)
+    return environ
+
+
+@pytest.fixture
+def mock_execve():
+    with mock.patch("os.execve", autospec=True) as mocked_function:
+        yield mocked_function
+
+
+def test_out_of_chroot(fake_env, mock_execve):
+    # When CROS_WORKON_SRCROOT is not set, we should not re-exec.
+    main.maybe_reexec(["--help"])
+    mock_execve.assert_not_called()
+
+
+def test_pythonpath_set(fake_env, mock_execve):
+    # With PYTHONPATH set, we should not re-exec.
+    fake_env["CROS_WORKON_SRCROOT"] = "/mnt/host/source"
+    fake_env["PYTHONPATH"] = "/foo/bar/baz"
+    main.maybe_reexec(["--help"])
+    mock_execve.assert_not_called()
+
+
+def test_zmake_does_not_exist(fake_env, mock_execve):
+    # When zmake is not at src/platform/ec/zephyr/zmake, don't re-exec.
+    fake_env["CROS_WORKON_SRCROOT"] = "/this/does/not/exist"
+    main.maybe_reexec(["--help"])
+    mock_execve.assert_not_called()
+
+
+def test_zmake_reexec(fake_env, mock_execve):
+    # Nothing else applies?  The re-exec should happen.
+    fake_env["CROS_WORKON_SRCROOT"] = "/mnt/host/source"
+    main.maybe_reexec(["--help"])
+    new_env = dict(fake_env)
+    new_env["PYTHONPATH"] = "/mnt/host/source/src/platform/ec/zephyr/zmake"
+    mock_execve.assert_called_once_with(
+        sys.executable,
+        [sys.executable, "-m", "zmake", "--help"],
+        new_env,
+    )
diff --git a/zephyr/zmake/zmake/__main__.py b/zephyr/zmake/zmake/__main__.py
index 74845d9bbc..a50ca3b08d 100644
--- a/zephyr/zmake/zmake/__main__.py
+++ b/zephyr/zmake/zmake/__main__.py
@@ -6,6 +6,7 @@
 import argparse
 import inspect
 import logging
+import os
 import pathlib
 import sys
 
@@ -13,6 +14,55 @@ import zmake.multiproc as multiproc
 import zmake.zmake as zm
 
 
+def maybe_reexec(argv):
+    """Re-exec zmake from the EC source tree, if possible and desired.
+
+    Zmake installs into the users' chroot, which makes it convenient
+    to execute, but can sometimes become tedious when zmake changes
+    land and users haven't upgraded their chroots yet.
+
+    We can partially subvert this problem by re-execing zmake from the
+    source if it's available.  This won't make it so developers never
+    need to upgrade their chroots (e.g., a toolchain upgrade could
+    require chroot upgrades), but at least makes it slightly more
+    convenient for an average repo sync.
+
+    Args:
+        argv: The argument list passed to the main function, not
+            including the executable path.
+
+    Returns:
+        None, if the re-exec did not happen, or never returns if the
+        re-exec did happen.
+    """
+    # We only re-exec if we are inside of a chroot (since if installed
+    # standalone using pip, there's already an "editable install"
+    # feature for that in pip.)
+    env = dict(os.environ)
+    srcroot = env.get("CROS_WORKON_SRCROOT")
+    if not srcroot:
+        return
+
+    # If for some reason we decide to move zmake in the future, then
+    # we don't want to use the re-exec logic.
+    zmake_path = (
+        pathlib.Path(srcroot) / "src" / "platform" / "ec" / "zephyr" / "zmake"
+    ).resolve()
+    if not zmake_path.is_dir():
+        return
+
+    # If PYTHONPATH is set, it is either because we just did a
+    # re-exec, or because the user wants to run a specific copy of
+    # zmake.  In either case, we don't want to re-exec.
+    if "PYTHONPATH" in env:
+        return
+
+    # Set PYTHONPATH so that we run zmake from source.
+    env["PYTHONPATH"] = str(zmake_path)
+
+    os.execve(sys.executable, [sys.executable, "-m", "zmake", *argv], env)
+
+
 def call_with_namespace(func, namespace):
     """Call a function with arguments applied from a Namespace.
 
@@ -55,6 +105,8 @@ def main(argv=None):
     if argv is None:
         argv = sys.argv[1:]
 
+    maybe_reexec(argv)
+
     parser = argparse.ArgumentParser()
     parser.add_argument(
         "--checkout", type=pathlib.Path, help="Path to ChromiumOS checkout"