Repository: pex-tool/pex

Issue:
Once a pyenv shim is activated and seen once by Pex it is always visible to Pex.
This happens because the target of a pyenv shims is mutable but the Pex interpreter cache is keyed by the shim (and its contents) which is immutable. Pex needs to factor in the mutable information (PYENV_VERSION, .python_version files in the CWD ancestor lineage and PYENV_ROOT/versions) into its cache key for pyenv shims.


Discussion:
The mutable information and its seach order is specified here: https://github.com/pyenv/pyenv#choosing-the-python-version

Resolving patch:
diff --git a/pex/common.py b/pex/common.py
--- a/pex/common.py
+++ b/pex/common.py
@@ -253,7 +253,7 @@ def register_rmtree(directory):
 
 
 def safe_mkdir(directory, clean=False):
-    # type: (str, bool) -> None
+    # type: (str, bool) -> str
     """Safely create a directory.
 
     Ensures a directory is present.  If it's not there, it is created.  If it is, it's a no-op. If
@@ -266,6 +266,8 @@ def safe_mkdir(directory, clean=False):
     except OSError as e:
         if e.errno != errno.EEXIST:
             raise
+    finally:
+        return directory
 
 
 def safe_open(filename, *args, **kwargs):
diff --git a/pex/interpreter.py b/pex/interpreter.py
--- a/pex/interpreter.py
+++ b/pex/interpreter.py
@@ -23,6 +23,7 @@
 from pex.jobs import ErrorHandler, Job, Retain, SpawnedJob, execute_parallel
 from pex.orderedset import OrderedSet
 from pex.platforms import Platform
+from pex.pyenv import Pyenv
 from pex.third_party.packaging import markers, tags
 from pex.third_party.pkg_resources import Distribution, Requirement
 from pex.tracer import TRACER
@@ -586,11 +587,55 @@ def _create_isolated_cmd(
 
         return cmd, env
 
+    # We use () as the unset sentinel for this lazily calculated cached value. The cached value
+    # itself should always be Optional[Pyenv].
+    #
+    # N.B.: The empty tuple type is not represented as Tuple[] as you might naivly guess but
+    # instead as Tuple[()].
+    #
+    # See:
+    # + https://github.com/python/mypy/issues/4211
+    # + https://www.python.org/dev/peps/pep-0484/#the-typing-module
+    _PYENV = ()  # type: Union[Tuple[()],Optional[Pyenv]]
+
+    @classmethod
+    def _pyenv(cls):
+        # type: () -> Optional[Pyenv]
+        if isinstance(cls._PYENV, tuple):
+            cls._PYENV = Pyenv.find()
+        return cls._PYENV
+
+    @classmethod
+    def _resolve_pyenv_shim(
+        cls,
+        binary,  # type: str
+        pyenv=None,  # type: Optional[Pyenv]
+    ):
+        # type: (...) -> Optional[str]
+
+        pyenv = pyenv or cls._pyenv()
+        if pyenv is not None:
+            shim = pyenv.as_shim(binary)
+            if shim is not None:
+                python = shim.select_version()
+                if python is None:
+                    TRACER.log("Detected inactive pyenv shim: {}.".format(shim), V=3)
+                else:
+                    TRACER.log("Detected pyenv shim activated to {}: {}.".format(python, shim), V=3)
+                return python
+        return binary
+
     INTERP_INFO_FILE = "INTERP-INFO"
 
     @classmethod
     def _spawn_from_binary_external(cls, binary):
-        def create_interpreter(stdout, check_binary=False):
+        # type: (str) -> SpawnedJob[PythonInterpreter]
+
+        def create_interpreter(
+            stdout,  # type: bytes
+            check_binary=False,  # type: bool
+        ):
+            # type: (...) -> PythonInterpreter
             identity = stdout.decode("utf-8").strip()
             if not identity:
                 raise cls.IdentificationError("Could not establish identity of {}.".format(binary))
@@ -707,6 +752,7 @@ def hashbang_matches(fn):
 
     @classmethod
     def _spawn_from_binary(cls, binary):
+        # type: (str) -> SpawnedJob[PythonInterpreter]
         canonicalized_binary = cls.canonicalize_path(binary)
         if not os.path.exists(canonicalized_binary):
             raise cls.InterpreterNotFound(canonicalized_binary)
@@ -721,14 +767,27 @@ def _spawn_from_binary(cls, binary):
         return cls._spawn_from_binary_external(canonicalized_binary)
 
     @classmethod
-    def from_binary(cls, binary):
-        # type: (str) -> PythonInterpreter
+    def from_binary(
+        cls,
+        binary,  # type: str
+        pyenv=None,  # type: Optional[Pyenv]
+    ):
+        # type: (...) -> PythonInterpreter
         """Create an interpreter from the given `binary`.
 
         :param binary: The path to the python interpreter binary.
+        :param pyenv: A custom Pyenv installation for handling pyenv shim identification.
+                      Auto-detected by default.
         :return: an interpreter created from the given `binary`.
         """
-        return cast(PythonInterpreter, cls._spawn_from_binary(binary).await_result())
+        python = cls._resolve_pyenv_shim(binary, pyenv=pyenv)
+        if python is None:
+            raise cls.IdentificationError("The pyenv shim at {} is not active.".format(binary))
+
+        try:
+            return cast(PythonInterpreter, cls._spawn_from_binary(python).await_result())
+        except Job.Error as e:
+            raise cls.IdentificationError("Failed to identify {}: {}".format(binary, e))
 
     @classmethod
     def _matches_binary_name(cls, path):
@@ -804,10 +863,12 @@ def iter_candidates():
             for path in cls._paths(paths=paths):
                 for fn in cls._expand_path(path):
                     if filter(fn):
-                        yield fn
+                        binary = cls._resolve_pyenv_shim(fn)
+                        if binary:
+                            yield binary
 
         results = execute_parallel(
-            inputs=list(iter_candidates()),
+            inputs=OrderedSet(iter_candidates()),
             spawn_func=cls._spawn_from_binary,
             error_handler=error_handler,
         )
diff --git a/pex/pyenv.py b/pex/pyenv.py
new file mode 100644
--- /dev/null
+++ b/pex/pyenv.py
@@ -0,0 +1,211 @@
+# Copyright 2021 Pants project contributors (see CONTRIBUTORS.md).
+# Licensed under the Apache License, Version 2.0 (see LICENSE).
+
+from __future__ import absolute_import
+
+import os
+import re
+import subprocess
+
+from pex.common import is_exe
+from pex.tracer import TRACER
+from pex.typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    import attr  # vendor:skip
+    from typing import Iterable, Iterator, Optional, Tuple, Iterable
+else:
+    from pex.third_party import attr
+
+
+@attr.s(frozen=True)
+class Pyenv(object):
+    root = attr.ib()  # type: str
+
+    @classmethod
+    def find(cls):
+        # type: () -> Optional[Pyenv]
+        """Finds the active pyenv installation if any."""
+        with TRACER.timed("Searching for pyenv root...", V=3):
+            pyenv_root = os.environ.get("PYENV_ROOT", "")
+            if not pyenv_root:
+                for path_entry in os.environ.get("PATH", "").split(os.pathsep):
+                    pyenv_exe = os.path.join(path_entry, "pyenv")
+                    if is_exe(pyenv_exe):
+                        process = subprocess.Popen(args=[pyenv_exe, "root"], stdout=subprocess.PIPE)
+                        stdout, _ = process.communicate()
+                        if process.returncode == 0:
+                            pyenv_root = str(stdout).strip()
+                            break
+
+            if pyenv_root:
+                pyenv = cls(pyenv_root)
+                TRACER.log("A pyenv installation was found: {}".format(pyenv), V=6)
+                return pyenv
+
+            TRACER.log("No pyenv installation was found.", V=6)
+            return None
+
+    @attr.s(frozen=True)
+    class Shim(object):
+        pyenv = attr.ib()  # type: Pyenv
+        path = attr.ib()  # type: str
+        name = attr.ib()  # type: str
+        major = attr.ib()  # type: Optional[str]
+        minor = attr.ib()  # type: Optional[str]
+
+        _SHIM_REGEX = re.compile(
+            r"""
+            ^
+            (?P<name>
+                python |
+                pypy
+            )
+            (?:
+                # Major version
+                (?P<major>[2-9])
+                (?:
+                    \.
+                    # Minor version
+                    (?P<minor>[0-9])
+                    # Some pyenv pythons include a suffix on the interpreter name, similar to
+                    # PEP-3149. For example, python3.6m to indicate it was built with pymalloc.
+                    [a-z]?
+                )?
+            )?
+            $
+            """,
+            flags=re.VERBOSE,
+        )
+
+        @classmethod
+        def parse(cls, pyenv, binary):
+            # type: (Pyenv, str) -> Optional[Pyenv.Shim]
+            """Parses shim information from a python binary path if it looks like a pyenv shim."""
+            if os.path.dirname(binary) != os.path.join(pyenv.root, "shims"):
+                return None
+            match = cls._SHIM_REGEX.match(os.path.basename(binary))
+            if match is None:
+                return None
+            return cls(
+                pyenv=pyenv,
+                path=binary,
+                name=match.group("name"),
+                major=match.group("major"),
+                minor=match.group("minor"),
+            )
+
+        _PYENV_CPYTHON_VERSION_LEADING_CHARS = frozenset(str(digit) for digit in range(2, 10))
+
+        def select_version(self, search_dir=None):
+            # type: (Optional[str]) -> Optional[str]
+            """Reports the active shim version for the given directory or $PWD.
+
+            If the shim is not activated, returns `None`.
+            """
+            with TRACER.timed("Calculating active version for {}...".format(self), V=6):
+                active_versions = self.pyenv.active_versions(search_dir=search_dir)
+                if active_versions:
+                    binary_name = os.path.basename(self.path)
+                    if self.name == "python" and not self.major and not self.minor:
+                        for pyenv_version in active_versions:
+                            if pyenv_version[0] in self._PYENV_CPYTHON_VERSION_LEADING_CHARS:
+                                TRACER.log(
+                                    "{} has active version {}".format(self, pyenv_version), V=6
+                                )
+                                return self.pyenv.python(pyenv_version, binary_name=binary_name)
+
+                    prefix = "{name}{major}{minor}".format(
+                        name="" if self.name == "python" else self.name,
+                        major=self.major or "",
+                        minor=".{}".format(self.minor) if self.minor else "",
+                    )
+                    for pyenv_version in active_versions:
+                        if pyenv_version.startswith(prefix):
+                            TRACER.log("{} has active version {}".format(self, pyenv_version), V=6)
+                            return self.pyenv.python(pyenv_version, binary_name=binary_name)
+
+                TRACER.log("{} is not activated.".format(self), V=6)
+                return None
+
+    def as_shim(self, binary):
+        # type: (str) -> Optional[Shim]
+        """View the given binary path as a pyenv shim script if it is one."""
+        return self.Shim.parse(self, binary)
+
+    @staticmethod
+    def _read_pyenv_versions(version_file):
+        # type: (str) -> Iterator[str]
+        with open(version_file) as fp:
+            for line in fp:
+                for version in line.strip().split():
+                    yield version
+
+    @staticmethod
+    def _find_local_version_file(search_dir):
+        # type: (str) -> Optional[str]
+        while True:
+            local_version_file = os.path.join(search_dir, ".python-version")
+            if os.path.exists(local_version_file):
+                return local_version_file
+            parent_dir = os.path.dirname(search_dir)
+            if parent_dir == search_dir:
+                return None
+            search_dir = parent_dir
+
+    def active_versions(self, search_dir=None):
+        # type: (Optional[str]) -> Tuple[str, ...]
+        """Reports the active pyenv versions for the given starting search directory or $PWD."""
+
+        source_and_versions = None  # type: Optional[Tuple[str, Iterable[str]]]
+
+        # See: https://github.com/pyenv/pyenv#choosing-the-python-version
+        with TRACER.timed("Finding {} active versions...".format(self), V=6):
+            shell_version = os.environ.get("PYENV_VERSION")
+            if shell_version:
+                source_and_versions = (
+                    "PYENV_VERSION={}".format(shell_version),
+                    shell_version.split(":"),
+                )
+            else:
+                cwd = search_dir if search_dir is not None else os.getcwd()
+                TRACER.log("Looking for pyenv version files starting from {}.".format(cwd), V=6)
+
+                local_version = self._find_local_version_file(search_dir=cwd)
+                if local_version:
+                    source_and_versions = (local_version, self._read_pyenv_versions(local_version))
+                else:
+                    global_version = os.path.join(self.root, "version")
+                    if os.path.exists(global_version):
+                        source_and_versions = (
+                            global_version,
+                            self._read_pyenv_versions(global_version),
+                        )
+
+        if source_and_versions:
+            source, versions = source_and_versions
+            TRACER.log("Found active versions in {}: {}".format(source, versions), V=6)
+            return tuple(versions)
+
+        TRACER.log("Found no active pyenv versions.", V=6)
+        return ()
+
+    def python(
+        self,
+        pyenv_version,  # type: str
+        binary_name=None,  # type: Optional[str]
+    ):
+        # type: (...) -> Optional[str]
+        """Return the path of the python binary for the given pyenv version.
+
+        Returns `None` if the given pyenv version is not installed.
+        """
+        # N.B.: Pyenv creates a 'python' symlink for both the CPython and PyPy versions it installs;
+        # so counting on 'python' is OK here. We do resolve the symlink though to return a canonical
+        # direct path to the python binary.
+        binary_name = binary_name or "python"
+
+        python = os.path.realpath(
+            os.path.join(self.root, "versions", pyenv_version, "bin", binary_name)
+        )
+        return python if is_exe(python) else None


Test patch:
diff --git a/pex/testing.py b/pex/testing.py
--- a/pex/testing.py
+++ b/pex/testing.py
@@ -517,15 +517,15 @@ def ensure_python_interpreter(version):
 
 @contextmanager
 def environment_as(**kwargs):
-    # type: (**str) -> Iterator[None]
+    # type: (**Any) -> Iterator[None]
     existing = {key: os.environ.get(key) for key in kwargs}
 
     def adjust_environment(mapping):
         for key, value in mapping.items():
             if value is not None:
-                os.environ[key] = value
+                os.environ[key] = str(value)
             else:
-                del os.environ[key]
+                os.environ.pop(key, None)
 
     adjust_environment(kwargs)
     try:
diff --git a/tests/test_interpreter.py b/tests/test_interpreter.py
--- a/tests/test_interpreter.py
+++ b/tests/test_interpreter.py
@@ -6,14 +6,16 @@
 import os
 import subprocess
 from collections import defaultdict
+from contextlib import contextmanager
 
 import pytest
 
 from pex import interpreter
-from pex.common import safe_mkdtemp, temporary_dir, touch
+from pex.common import safe_mkdir, safe_mkdtemp, temporary_dir, touch
 from pex.compatibility import PY3
 from pex.executor import Executor
 from pex.interpreter import PythonInterpreter
+from pex.pyenv import Pyenv
 from pex.testing import (
     PY27,
     PY35,
@@ -34,7 +36,7 @@
     from unittest.mock import Mock, patch  # type: ignore[misc,no-redef,import]
 
 if TYPE_CHECKING:
-    from typing import Any, Iterator, List, Tuple
+    from typing import Any, Iterator, List, Tuple, Optional
 
     from pex.interpreter import InterpreterOrError
 
@@ -203,8 +205,8 @@ def test_iter_interpreter_path_filter_symlink(self, test_interpreter1, test_inte
             )
             assert os.path.basename(expected_interpreter.binary) != "jake"
 
-    def test_pyenv_shims(self):
-        # type: () -> None
+    def test_pyenv_shims(self, tmpdir):
+        # type: (Any) -> None
         py35, _, run_pyenv = ensure_python_distribution(PY35)
         py36 = ensure_python_interpreter(PY36)
 
@@ -212,31 +214,89 @@ def test_pyenv_shims(self):
         pyenv_shims = os.path.join(pyenv_root, "shims")
 
         def pyenv_global(*versions):
+            # type: (*str) -> None
             run_pyenv(["global"] + list(versions))
 
-        def assert_shim(shim_name, expected_binary_path):
-            python = PythonInterpreter.from_binary(os.path.join(pyenv_shims, shim_name))
-            assert expected_binary_path == python.binary
-
-        with temporary_dir() as pex_root:
-            with ENV.patch(PEX_ROOT=pex_root) as pex_env:
-                with environment_as(PYENV_ROOT=pyenv_root, **pex_env):
-                    pyenv_global(PY35, PY36)
-                    assert_shim("python3", py35)
-
-                    pyenv_global(PY36, PY35)
-                    # The python3 shim is now pointing at python3.6 but the Pex cache has a valid
-                    # entry for the old python3.5 association (the interpreter still exists.)
-                    assert_shim("python3", py35)
-
-                    # The shim pointer is now invalid since python3.5 was uninstalled and so should
-                    # be re-read.
-                    py35_deleted = "{}.uninstalled".format(py35)
-                    os.rename(py35, py35_deleted)
-                    try:
-                        assert_shim("python3", py36)
-                    finally:
-                        os.rename(py35_deleted, py35)
+        def pyenv_local(*versions):
+            # type: (*str) -> None
+            run_pyenv(["local"] + list(versions))
+
+        @contextmanager
+        def pyenv_shell(*versions):
+            # type: (*str) -> Iterator[None]
+            with environment_as(PYENV_VERSION=":".join(versions)):
+                yield
+
+        pex_root = os.path.join(str(tmpdir), "pex_root")
+        cwd = safe_mkdir(os.path.join(str(tmpdir), "home", "jake", "project"))
+        with ENV.patch(PEX_ROOT=pex_root) as pex_env, environment_as(
+            PYENV_ROOT=pyenv_root, PEX_PYTHON_PATH=pyenv_shims, **pex_env
+        ), pyenv_shell(), pushd(cwd):
+            pyenv = Pyenv.find()
+            assert pyenv is not None
+            assert pyenv_root == pyenv.root
+
+            def interpreter_for_shim(shim_name):
+                # type: (str) -> PythonInterpreter
+                binary = os.path.join(pyenv_shims, shim_name)
+                return PythonInterpreter.from_binary(binary, pyenv=pyenv)
+
+            def assert_shim(
+                shim_name,  # type: str
+                expected_binary_path,  # type: str
+            ):
+                # type: (...) -> None
+                python = interpreter_for_shim(shim_name)
+                assert expected_binary_path == python.binary
+
+            def assert_shim_inactive(shim_name):
+                # type: (str) -> None
+                with pytest.raises(PythonInterpreter.IdentificationError):
+                    interpreter_for_shim(shim_name)
+
+            pyenv_global(PY35, PY36)
+            assert_shim("python", py35)
+            assert_shim("python3", py35)
+            assert_shim("python3.5", py35)
+            assert_shim("python3.6", py36)
+
+            pyenv_global(PY36, PY35)
+            assert_shim("python", py36)
+            assert_shim("python3", py36)
+            assert_shim("python3.6", py36)
+            assert_shim("python3.5", py35)
+
+            pyenv_local(PY35)
+            assert_shim("python", py35)
+            assert_shim("python3", py35)
+            assert_shim("python3.5", py35)
+            assert_shim_inactive("python3.6")
+
+            with pyenv_shell(PY36):
+                assert_shim("python", py36)
+                assert_shim("python3", py36)
+                assert_shim("python3.6", py36)
+                assert_shim_inactive("python3.5")
+
+            with pyenv_shell(PY35, PY36):
+                assert_shim("python", py35)
+                assert_shim("python3", py35)
+                assert_shim("python3.5", py35)
+                assert_shim("python3.6", py36)
+
+            # The shim pointer is now invalid since python3.5 was uninstalled and so
+            # should be re-read and found invalid.
+            py35_version_dir = os.path.dirname(os.path.dirname(py35))
+            py35_deleted = "{}.uninstalled".format(py35_version_dir)
+            os.rename(py35_version_dir, py35_deleted)
+            try:
+                assert_shim_inactive("python")
+                assert_shim_inactive("python3")
+                assert_shim_inactive("python3.5")
+            finally:
+                os.rename(py35_deleted, py35_version_dir)
+
+            assert_shim("python", py35)
 
 
 def test_latest_release_of_min_compatible_version():