Repository: pex-tool/pex

Issue:
When Pex is run from a Pex PEX its isolation is broken.
As noted here: https://github.com/pantsbuild/pex/commit/4562d7d386a56095612b6f7f9b8eb8261e45c1af#r47156419


Discussion:

Resolving patch:
diff --git a/pex/third_party/__init__.py b/pex/third_party/__init__.py
--- a/pex/third_party/__init__.py
+++ b/pex/third_party/__init__.py
@@ -375,7 +375,18 @@ def recursive_copy(srcdir, dstdir):
 
         pex_path = os.path.join(vendor.VendorSpec.ROOT, "pex")
         with _tracer().timed("Hashing pex"):
-            dir_hash = CacheHelper.dir_hash(pex_path)
+            if os.path.isdir(pex_path):
+                dir_hash = CacheHelper.dir_hash(pex_path)
+            else:
+                pex_zip = os.path.abspath(sys.argv[0])
+                assert zipfile.is_zipfile(pex_zip) and pex_zip == os.path.commonprefix(
+                    (pex_zip, pex_path)
+                ), (
+                    "Expected the `pex` module to be available via an installed distribution or "
+                    "else via a PEX zipfile present as argv0. Loaded the `pex` module from {} and "
+                    "argv0 is {}.".format(pex_path, sys.argv[0])
+                )
+                dir_hash = CacheHelper.zip_hash(pex_zip, os.path.relpath(pex_path, pex_zip))
         isolated_dir = os.path.join(ENV.PEX_ROOT, "isolated", dir_hash)
 
         with _tracer().timed("Isolating pex"):
diff --git a/pex/util.py b/pex/util.py
--- a/pex/util.py
+++ b/pex/util.py
@@ -11,7 +11,14 @@
 from site import makepath  # type: ignore[attr-defined]
 from zipfile import ZipFile
 
-from pex.common import atomic_directory, filter_pyc_dirs, filter_pyc_files, safe_mkdir, safe_mkdtemp
+from pex.common import (
+    atomic_directory,
+    filter_pyc_dirs,
+    filter_pyc_files,
+    open_zip,
+    safe_mkdir,
+    safe_mkdtemp,
+)
 from pex.compatibility import (  # type: ignore[attr-defined]  # `exec_function` is defined dynamically
     PY2,
     exec_function,
@@ -24,14 +31,14 @@
     resource_string,
 )
 from pex.tracer import TRACER
-from pex.typing import TYPE_CHECKING
+from pex.typing import TYPE_CHECKING, cast
 
 if TYPE_CHECKING:
     if PY2:
         from hashlib import _hash as _Hash
     else:
         from hashlib import _Hash
-    from typing import Any, Callable, IO, Iterable, Iterator, Optional
+    from typing import Any, BinaryIO, Callable, IO, Iterable, Iterator, Optional
 
 
 class DistributionHelper(object):
@@ -104,7 +111,7 @@ def distribution_from_path(cls, path, name=None):
 class CacheHelper(object):
     @classmethod
     def update_hash(cls, filelike, digest):
-        # type: (IO[bytes], _Hash) -> None
+        # type: (BinaryIO, _Hash) -> None
         """Update the digest of a single file in a memory-efficient manner."""
         block_size = digest.block_size * 1024
         for chunk in iter(lambda: filelike.read(block_size), b""):
@@ -122,7 +129,7 @@ def hash(cls, path, digest=None, hasher=sha1):
 
     @classmethod
     def _compute_hash(cls, names, stream_factory):
-        # type: (Iterable[str], Callable[[str], IO]) -> str
+        # type: (Iterable[str], Callable[[str], BinaryIO]) -> str
         digest = sha1()
         # Always use / as the path separator, since that's what zip uses.
         hashed_names = [n.replace(os.sep, "/") for n in names]
@@ -148,8 +155,8 @@ def pex_code_hash(cls, d):
         names = sorted(f for f in cls._iter_non_pyc_files(d) if not f.startswith("."))
 
         def stream_factory(name):
-            # type: (str) -> IO
-            return open(os.path.join(d, name), "rb")  # noqa: T802
+            # type: (str) -> BinaryIO
+            return cast("BinaryIO", open(os.path.join(d, name), "rb"))
 
         return cls._compute_hash(names, stream_factory)
 
@@ -160,11 +167,34 @@ def dir_hash(cls, d):
         names = sorted(cls._iter_non_pyc_files(d))
 
         def stream_factory(name):
-            # type: (str) -> IO
-            return open(os.path.join(d, name), "rb")  # noqa: T802
+            # type: (str) -> BinaryIO
+            return cast("BinaryIO", open(os.path.join(d, name), "rb"))
 
         return cls._compute_hash(names, stream_factory)
 
+    @classmethod
+    def zip_hash(
+        cls,
+        zip_file,  # type: str
+        relpath=None,  # type: Optional[str]
+    ):
+        # type: (...) -> str
+        """Return a reproducible hash of the contents of a zip; excluding all `.pyc` files."""
+        with open_zip(zip_file) as zf:
+            names = sorted(
+                filter_pyc_files(
+                    name
+                    for name in zf.namelist()
+                    if not name.endswith("/") and not relpath or name.startswith(relpath)
+                )
+            )
+
+            def stream_factory(name):
+                # type: (str) -> BinaryIO
+                return cast("BinaryIO", zf.open(name, "r"))
+
+            return cls._compute_hash(names, stream_factory)
+
     @classmethod
     def cache_distribution(cls, zf, source, target_dir):
         # type: (ZipFile, str, str) -> Distribution


Test patch:
diff --git a/tests/test_integration.py b/tests/test_integration.py
--- a/tests/test_integration.py
+++ b/tests/test_integration.py
@@ -71,6 +71,7 @@
         Callable,
         ContextManager,
         Dict,
+        FrozenSet,
         Iterable,
         Iterator,
         List,
@@ -3006,3 +3007,158 @@ def test_2020_resolver_engaged_issues_1179():
     assert "1: boto3==1.15.6 requires botocore<1.19.0,>=1.18.6 but " in results.error
 
     run_pex_command(args=["--resolver-version", "pip-2020-resolver"] + pex_args).assert_success()
+
+
+def test_isolated_pex_zip_issues_1232(tmpdir):
+    # type: (Any) -> None
+
+    pex_root = os.path.join(str(tmpdir), "pex_root")
+
+    python35 = ensure_python_interpreter(PY35)
+    python36 = ensure_python_interpreter(PY36)
+
+    pex_env = make_env(PEX_PYTHON_PATH=os.pathsep.join((python35, python36)))
+
+    def add_pex_args(*args):
+        # type: (*str) -> List[str]
+        return list(args) + [
+            "--pex-root",
+            pex_root,
+            "--runtime-pex-root",
+            pex_root,
+            "--interpreter-constraint",
+            "CPython=={version}".format(version=PY35),
+        ]
+
+    def tally_isolated_vendoreds():
+        # type: () -> Dict[str, FrozenSet[str]]
+        def vendored_toplevel(isolated_dir):
+            # type: (str) -> Iterator[str]
+            vendored_dir = os.path.join(isolated_dir, "pex/vendor/_vendored")
+            for path in os.listdir(vendored_dir):
+                if path in ("__pycache__", "__init__.py"):
+                    continue
+                if os.path.isdir(os.path.join(vendored_dir, path)):
+                    yield path
+                module, ext = os.path.splitext(path)
+                if ext == ".py":
+                    yield module
+
+        isolated_root = os.path.join(pex_root, "isolated")
+        vendored_by_isolated = {}
+        for entry in os.listdir(isolated_root):
+            path = os.path.join(isolated_root, entry)
+            if not os.path.isdir(path):
+                continue
+            vendored_by_isolated[path] = frozenset(vendored_toplevel(path))
+        return vendored_by_isolated
+
+    # 1. Isolate current loose source Pex at build-time.
+    # ===
+    current_pex_pex = os.path.join(str(tmpdir), "pex-current.pex")
+    results = run_pex_command(
+        args=add_pex_args(".", "-c", "pex", "-o", current_pex_pex), env=pex_env, python=python35
+    )
+    results.assert_success()
+
+    current_isolated_vendoreds = tally_isolated_vendoreds()
+    assert 1 == len(current_isolated_vendoreds), (
+        "Since we just ran the Pex tool and nothing else, a single isolation of the Pex loose "
+        "source in this repo should have occurred."
+    )
+    assert {"pip", "wheel"}.issubset(
+        list(current_isolated_vendoreds.values())[0]
+    ), "Expected isolation of current Pex code to be a full build-time isolation."
+
+    # 2. Isolate current Pex PEX at run-time.
+    # ===
+    modified_pex_src = os.path.join(str(tmpdir), "modified_pex_src")
+    shutil.copytree("pex", os.path.join(modified_pex_src, "pex"))
+    with open(os.path.join(modified_pex_src, "pex", "version.py"), "a") as fp:
+        fp.write("# modified\n")
+    shutil.copy("pyproject.toml", os.path.join(modified_pex_src, "pyproject.toml"))
+    # N.B.: README.rst is needed by flit since we tell it to pull the distribution description from
+    # there when building the Pex distribution.
+    shutil.copy("README.rst", os.path.join(modified_pex_src, "README.rst"))
+
+    modified_pex = os.path.join(str(tmpdir), "modified.pex")
+    subprocess.check_call(
+        args=add_pex_args(
+            python36, current_pex_pex, modified_pex_src, "-c", "pex", "-o", modified_pex
+        ),
+        env=pex_env,
+    )
+    current_pex_isolated_vendoreds = tally_isolated_vendoreds()
+    current_pex_isolation = set(current_isolated_vendoreds.keys()) ^ set(
+        current_pex_isolated_vendoreds.keys()
+    )
+    assert 1 == len(current_pex_isolation), (
+        "Since the modified Pex PEX was built from a Pex PEX an isolation of the Pex PEX bootstrap "
+        "code should have occurred bringing the total isolations up to two."
+    )
+    current_pex_vendoreds = current_pex_isolated_vendoreds[current_pex_isolation.pop()]
+    assert "pip" not in current_pex_vendoreds, "Expected a Pex runtime isolation."
+    assert "wheel" not in current_pex_vendoreds, "Expected a Pex runtime isolation."
+
+    # 3. Isolate modified Pex PEX at build-time.
+    # ===
+    ansicolors_pex = os.path.join(str(tmpdir), "ansicolors.pex")
+    subprocess.check_call(
+        args=add_pex_args(
+            python36,
+            modified_pex,
+            "ansicolors==1.1.8",
+            "-o",
+            ansicolors_pex,
+        ),
+        env=pex_env,
+    )
+    modified_pex_isolated_vendoreds = tally_isolated_vendoreds()
+    modified_pex_isolation = set(current_pex_isolated_vendoreds.keys()) ^ set(
+        modified_pex_isolated_vendoreds.keys()
+    )
+    assert 1 == len(modified_pex_isolation), (
+        "Since the ansicolors PEX was built from the modifed Pex PEX a new isolation of the "
+        "modified Pex PEX code should have occurred bringing the total isolations up to three."
+    )
+    assert {"pip", "wheel"}.issubset(
+        modified_pex_isolated_vendoreds[modified_pex_isolation.pop()]
+    ), "Expected isolation of modified Pex code to be a full build-time isolation."
+
+    # 4. Isolate modified Pex PEX at run-time.
+    # ===
+    # Force the bootstrap to run interpreter identification which will force a Pex isolation.
+    shutil.rmtree(os.path.join(pex_root, "interpreters"))
+    subprocess.check_call(args=[python36, ansicolors_pex, "-c", "import colors"], env=pex_env)
+    ansicolors_pex_isolated_vendoreds = tally_isolated_vendoreds()
+    ansicolors_pex_isolation = set(modified_pex_isolated_vendoreds.keys()) ^ set(
+        ansicolors_pex_isolated_vendoreds.keys()
+    )
+    assert 1 == len(ansicolors_pex_isolation), (
+        "Since the ansicolors PEX has modified Pex bootstrap code, a further isolation should have"
+        "occurred bringing the total isolations up to four."
+    )
+    ansicolors_pex_vendoreds = ansicolors_pex_isolated_vendoreds[ansicolors_pex_isolation.pop()]
+    assert "pip" not in ansicolors_pex_vendoreds, "Expected a Pex runtime isolation."
+    assert "wheel" not in ansicolors_pex_vendoreds, "Expected a Pex runtime isolation."
+
+    # 5. No new isolations.
+    # ===
+    ansicolors_pex = os.path.join(str(tmpdir), "ansicolors.old.pex")
+    subprocess.check_call(
+        args=add_pex_args(
+            python36,
+            modified_pex,
+            "ansicolors==1.0.2",
+            "-o",
+            ansicolors_pex,
+        ),
+        env=pex_env,
+    )
+
+    # Force the bootstrap to run interpreter identification which will force a Pex isolation.
+    shutil.rmtree(os.path.join(pex_root, "interpreters"))
+    subprocess.check_call(args=[python36, ansicolors_pex, "-c", "import colors"], env=pex_env)
+    assert (
+        ansicolors_pex_isolated_vendoreds == tally_isolated_vendoreds()
+    ), "Expecting no new Pex isolations."