Repository: pex-tool/pex

Issue:
Pex `--venv` mode does not account for `PEX_PATH`.
When `PEX_PATH` is specified at runtime, it should trigger a venv re-build as needed.
Currently we only rebuild for interpreter contraints changes, `PEX_PYTHON` and `PEX_PYTHON_PATH` changes; i.e.: changes that could affect the correct interpreter to build the venv with:
https://github.com/pantsbuild/pex/blob/dd5c86f38ef76090089f88da8a4cefd8769a47c9/pex/variables.py#L648-L668


Discussion:

Resolving patch:
diff --git a/pex/pex_builder.py b/pex/pex_builder.py
--- a/pex/pex_builder.py
+++ b/pex/pex_builder.py
@@ -86,7 +86,9 @@ def __maybe_run_venv__(pex):
   from pex.tracer import TRACER
   from pex.variables import venv_dir
 
-  venv_home = venv_dir({pex_root!r}, {pex_hash!r}, {interpreter_constraints!r})
+  venv_home = venv_dir(
+    {pex_root!r}, {pex_hash!r}, {interpreter_constraints!r}, pex_path={pex_path!r}
+  )
   venv_pex = os.path.join(venv_home, 'pex')
   if not is_exe(venv_pex):
     # Code in bootstrap_pex will (re)create the venv after selecting the correct interpreter. 
@@ -512,6 +514,7 @@ def _prepare_code(self):
             pex_root=self._pex_info.raw_pex_root,
             pex_hash=self._pex_info.pex_hash,
             interpreter_constraints=self._pex_info.interpreter_constraints,
+            pex_path=self._pex_info.pex_path,
             is_unzip=self._pex_info.unzip,
             is_venv=self._pex_info.venv,
         )
diff --git a/pex/pex_info.py b/pex/pex_info.py
--- a/pex/pex_info.py
+++ b/pex/pex_info.py
@@ -257,7 +257,9 @@ def venv_dir(self):
             return None
         if self.pex_hash is None:
             raise ValueError("The venv_dir was requested but no pex_hash was set.")
-        return variables.venv_dir(self.pex_root, self.pex_hash, self.interpreter_constraints)
+        return variables.venv_dir(
+            self.pex_root, self.pex_hash, self.interpreter_constraints, pex_path=self.pex_path
+        )
 
     @property
     def strip_pex_env(self):
diff --git a/pex/variables.py b/pex/variables.py
--- a/pex/variables.py
+++ b/pex/variables.py
@@ -7,6 +7,7 @@
 from __future__ import absolute_import
 
 import hashlib
+import json
 import os
 import sys
 from contextlib import contextmanager
@@ -18,7 +19,18 @@
 from pex.venv_bin_path import BinPath
 
 if TYPE_CHECKING:
-    from typing import Callable, Dict, Iterable, Iterator, Optional, Tuple, TypeVar, Type, Union
+    from typing import (
+        Callable,
+        Dict,
+        Iterable,
+        Iterator,
+        Optional,
+        Tuple,
+        TypeVar,
+        Type,
+        Union,
+        Any,
+    )
 
     _O = TypeVar("_O")
     _P = TypeVar("_P")
@@ -649,20 +661,40 @@ def venv_dir(
     pex_root,  # type: str
     pex_hash,  # type: str
     interpreter_constraints,  # type: Iterable[str]
+    pex_path=None,  # type: Optional[str]
 ):
     # type: (...) -> str
-    hasher = hashlib.sha1()
-    hasher.update(
-        "interpreter_constraints:{}".format(" or ".join(sorted(interpreter_constraints))).encode(
-            "utf-8"
-        )
-    )
-    hasher.update("PEX_PYTHON:{}".format(ENV.PEX_PYTHON).encode("utf-8"))
-    hasher.update("PEX_PYTHON_PATH:{}".format(ENV.PEX_PYTHON_PATH).encode("utf-8"))
-    interpreter_selection_hash = hasher.hexdigest()
+    # The venv contents are affected by which interpreter is selected as well as which PEX files are
+    # in play. The former can be influenced by interpreter constraints changes as well as PEX_PYTHON
+    # and PEX_PYTHON_PATH. The latter is influenced via PEX_PATH.
+    venv_contents = {
+        "interpreter_constraints": sorted(interpreter_constraints),
+        "PEX_PYTHON": ENV.PEX_PYTHON,
+        "PEX_PYTHON_PATH": ENV.PEX_PYTHON_PATH,
+        "pex_path": {},
+    }  # type: Dict[str, Any]
+
+    # PexInfo.pex_path and PEX_PATH are merged by PEX at runtime, so we include both and hash just
+    # the distributions since those are the only items used from PEX_PATH adjoined PEX files; i.e.:
+    # neither the entry_point nor any other PEX file data or metadata is used.
+    def add_pex_path_items(pex_path):
+        if not pex_path:
+            return
+        from pex.pex_info import PexInfo
+
+        pex_path_contents = venv_contents["pex_path"]
+        for pex in pex_path.split(":"):
+            pex_path_contents[pex] = PexInfo.from_pex(pex).distributions
+
+    add_pex_path_items(pex_path)
+    add_pex_path_items(ENV.PEX_PATH)
+
+    venv_contents_hash = hashlib.sha1(
+        json.dumps(venv_contents, sort_keys=True).encode("utf-8")
+    ).hexdigest()
     return os.path.join(
         _expand_pex_root(pex_root),
         "venvs",
         pex_hash,
-        interpreter_selection_hash,
+        venv_contents_hash,
     )


Test patch:
diff --git a/tests/test_integration.py b/tests/test_integration.py
--- a/tests/test_integration.py
+++ b/tests/test_integration.py
@@ -2746,6 +2746,80 @@ def get_fabric_versions(pex):
     assert invoke_pex_info.pex_hash != fabric_pex_info.pex_hash
 
 
+def test_venv_mode_pex_path_issues_1225(tmpdir):
+    # type: (Any) -> None
+
+    test_file = os.path.join(str(tmpdir), "test.py")
+    with open(test_file, "w") as fp:
+        fp.write(
+            dedent(
+                """
+                import sys
+
+                try:
+                    __import__(sys.argv[1])
+                except ImportError:
+                    sys.exit(int(sys.argv[2]))
+                """
+            )
+        )
+
+    empty_pex = os.path.join(str(tmpdir), "empty.pex")
+    results = run_pex_command(args=["--venv", "-o", empty_pex])
+    results.assert_success()
+
+    output, returncode = run_simple_pex(empty_pex, args=[test_file, "colors", "37"])
+    assert 37 == returncode, output.decode("utf-8")
+
+    colors_pex = os.path.join(str(tmpdir), "colors.pex")
+    results = run_pex_command(args=["ansicolors==1.1.8", "-o", colors_pex])
+    results.assert_success()
+
+    # Exporting PEX_PATH should re-create the venv.
+    output, returncode = run_simple_pex(
+        empty_pex, args=[test_file, "colors", "37"], env=make_env(PEX_PATH=colors_pex)
+    )
+    assert 0 == returncode, output.decode("utf-8")
+
+    results = run_pex_command(args=["--pex-path", colors_pex, "--venv", "-o", empty_pex])
+    results.assert_success()
+
+    output, returncode = run_simple_pex(empty_pex, args=[test_file, "colors", "37"])
+    assert 0 == returncode
+
+    # Exporting PEX_PATH should re-create the venv, adding to --pex-path.
+    pkginfo_pex = os.path.join(str(tmpdir), "pkginfo.pex")
+    results = run_pex_command(args=["pkginfo==1.7.0", "-o", pkginfo_pex])
+    results.assert_success()
+
+    pex_path_env = make_env(PEX_PATH=pkginfo_pex)
+    output, returncode = run_simple_pex(
+        empty_pex, args=[test_file, "colors", "37"], env=pex_path_env
+    )
+    assert 0 == returncode
+    output, returncode = run_simple_pex(
+        empty_pex, args=[test_file, "pkginfo", "42"], env=pex_path_env
+    )
+    assert 0 == returncode
+
+    # Exporting PEX_PATH should re-create the venv since the adjoined pex file's distribution
+    # contents have changed.
+    results = run_pex_command(args=["ascii-ruler==0.0.4", "-o", pkginfo_pex])
+    results.assert_success()
+    output, returncode = run_simple_pex(
+        empty_pex, args=[test_file, "colors", "37"], env=pex_path_env
+    )
+    assert 0 == returncode
+    output, returncode = run_simple_pex(
+        empty_pex, args=[test_file, "ascii_ruler", "19"], env=pex_path_env
+    )
+    assert 0 == returncode
+    output, returncode = run_simple_pex(
+        empty_pex, args=[test_file, "pkginfo", "42"], env=pex_path_env
+    )
+    assert 42 == returncode
+
+
 @pytest.mark.parametrize(
     "mode_args",
     [