Fix regression in venv symlinking. (#2090)

The change in #2033 that safeguarded loose `--venv` PEXes such that they
could be executed, moved and continue to execute properly from their new
location, silently broke zipapp and packed layout `--venv` PEX
symlinking. Although `--venv` PEXes for those layouts continued to work,
they used copying to build their venvs instead of symlinking even when
symlinking (the default for `--venv` PEXes) was the intended style.

Fixes #2088

pex/layout.py

@@ -21,11 +21,24 @@
 BOOTSTRAP_DIR = ".bootstrap"
 DEPS_DIR = ".deps"
 PEX_INFO_PATH = "PEX-INFO"
+PEX_LAYOUT_PATH = "PEX-LAYOUT"
 
 
 class Layout(Enum["Layout.Value"]):
     class Value(Enum.Value):
-        pass
+        @classmethod
+        def try_load(cls, pex_directory):
+            # type: (str) -> Optional[Layout.Value]
+            layout = os.path.join(pex_directory, PEX_LAYOUT_PATH)
+            if not os.path.isfile(layout):
+                return None
+            with open(layout) as fp:
+                return Layout.for_value(fp.read().strip())
+
+        def record(self, pex_directory):
+            # type: (str) -> None
+            with open(os.path.join(pex_directory, PEX_LAYOUT_PATH), "w") as fp:
+                fp.write(self.value)
 
     ZIPAPP = Value("zipapp")
     PACKED = Value("packed")
@@ -47,10 +60,23 @@ def identify(cls, pex):
 
         return cls.LOOSE
 
+    @classmethod
+    def identify_original(cls, pex):
+        # type: (str) -> Layout.Value
+        layout = cls.identify(pex)
+        if layout != Layout.LOOSE:
+            return layout
+        return cls.Value.try_load(pex) or Layout.LOOSE
+
 
 class _Layout(object):
-    def __init__(self, path):
-        # type: (str) -> None
+    def __init__(
+        self,
+        layout,  # type: Layout.Value
+        path,  # type: str
+    ):
+        # type: (...) -> None
+        self._layout = layout
         self._path = os.path.normpath(path)
 
     @property
@@ -95,6 +121,10 @@ def extract_main(self, dest_dir):
         # type: (str) -> None
         raise NotImplementedError()
 
+    def record(self, dest_dir):
+        # type: (str) -> None
+        self._layout.record(dest_dir)
+
 
 def _install(
     layout,  # type: _Layout
@@ -165,7 +195,7 @@ def _install(
 
                     layout.extract_pex_info(chroot.work_dir)
                     layout.extract_main(chroot.work_dir)
-
+                    layout.record(chroot.work_dir)
         return install_to
 
 
@@ -176,7 +206,7 @@ def __init__(
         zfp,  # type: zipfile.ZipFile
     ):
         # type: (...) -> None
-        super(_ZipAppPEX, self).__init__(path)
+        super(_ZipAppPEX, self).__init__(Layout.ZIPAPP, path)
         self._zfp = zfp
         self._names = tuple(zfp.namelist())
 
@@ -224,6 +254,10 @@ def __str__(self):
 
 
 class _PackedPEX(_Layout):
+    def __init__(self, path):
+        # type: (str) -> None
+        super(_PackedPEX, self).__init__(Layout.PACKED, path)
+
     def extract_bootstrap(self, dest_dir):
         # type: (str) -> None
         with open_zip(os.path.join(self._path, BOOTSTRAP_DIR)) as zfp:

pex/pex.py

@@ -171,7 +171,7 @@ def __init__(
     def layout(self):
         # type: () -> Layout.Value
         if self._layout is None:
-            self._layout = Layout.identify(self._pex)
+            self._layout = Layout.identify_original(self._pex)
         return self._layout
 
     def pex_info(self, include_env_overrides=True):

pex/venv/pex.py

@@ -10,7 +10,7 @@
 from collections import Counter, defaultdict
 from textwrap import dedent
 
-from pex import pex_warnings
+from pex import layout, pex_warnings
 from pex.common import chmod_plus_x, pluralize, safe_mkdir
 from pex.compatibility import is_valid_python_identifier
 from pex.dist_metadata import Distribution
@@ -313,11 +313,12 @@ def _populate_sources(
         src=PEXEnvironment.mount(pex.path()).path,
         dst=venv.site_packages_dir,
         exclude=(
-            pex_info.internal_cache,
-            pex_info.bootstrap,
             "__main__.py",
             "__pycache__",
-            pex_info.PATH,
+            layout.BOOTSTRAP_DIR,
+            layout.DEPS_DIR,
+            layout.PEX_INFO_PATH,
+            layout.PEX_LAYOUT_PATH,
         ),
         symlink=False,
     ):

tests/integration/test_issue_2088.py

@@ -0,0 +1,90 @@
+# Copyright 2023 Pants project contributors (see CONTRIBUTORS.md).
+# Licensed under the Apache License, Version 2.0 (see LICENSE).
+import os.path
+import shutil
+import subprocess
+import sys
+
+import pytest
+
+from pex.common import safe_rmtree
+from pex.compatibility import commonpath
+from pex.layout import Layout
+from pex.testing import run_pex_command
+from pex.typing import TYPE_CHECKING
+from pex.venv.virtualenv import Virtualenv
+
+if TYPE_CHECKING:
+    from typing import Any
+
+
+@pytest.mark.parametrize(
+    "layout", [pytest.param(layout, id=layout.value) for layout in Layout.values()]
+)
+@pytest.mark.parametrize(
+    "symlink_site_packages",
+    [
+        pytest.param(True, id="--no-venv-site-packages-copies"),
+        pytest.param(False, id="--venv-site-packages-copies"),
+    ],
+)
+def test_venv_symlink_site_packages(
+    tmpdir,  # type: Any
+    layout,  # type: Layout.Value
+    symlink_site_packages,  # type: bool
+):
+    # type: (...) -> None
+
+    pex = os.path.join(str(tmpdir), "pex")
+    pex_root = os.path.join(str(tmpdir), "pex_root")
+    venv_site_packages_copies_arg = (
+        "--no-venv-site-packages-copies" if symlink_site_packages else "--venv-site-packages-copies"
+    )
+    result = run_pex_command(
+        args=[
+            "ansicolors==1.1.8",
+            "--pex-root",
+            pex_root,
+            "--runtime-pex-root",
+            pex_root,
+            "--venv",
+            venv_site_packages_copies_arg,
+            "--layout",
+            layout.value,
+            "--seed",
+            "-o",
+            pex,
+        ]
+    )
+    result.assert_success()
+    venv_pex_path = str(result.output.strip())
+
+    safe_rmtree(pex_root)
+    colors_module_realpath = os.path.realpath(
+        subprocess.check_output(
+            args=[sys.executable, pex, "-c", "import colors; print(colors.__file__)"]
+        )
+        .decode("utf-8")
+        .strip()
+    )
+
+    venv_dir = os.path.dirname(venv_pex_path)
+    virtualenv = Virtualenv(venv_dir)
+    site_packages_dir = os.path.realpath(virtualenv.site_packages_dir)
+
+    ansicolors_venv_package_dir_realpath = os.path.join(site_packages_dir, "colors")
+    assert os.path.isdir(ansicolors_venv_package_dir_realpath)
+
+    symlinks_expected = symlink_site_packages and layout != Layout.LOOSE
+    assert os.path.islink(ansicolors_venv_package_dir_realpath) == symlinks_expected
+
+    if symlinks_expected:
+        installed_wheels_dir_realpath = os.path.realpath(os.path.join(pex_root, "installed_wheels"))
+        assert installed_wheels_dir_realpath == commonpath(
+            (installed_wheels_dir_realpath, colors_module_realpath)
+        )
+    else:
+        assert (
+            os.path.join(ansicolors_venv_package_dir_realpath, "__init__.py")
+            == colors_module_realpath
+        )