Repository: pex-tool/pex

Issue:
The `--venv-site-packages-copies` mode should merge directories.
As a simplification in code, when the `--venv-site-packages-copies` mode was introduced to allow toggling `--venv` mode PEXes between symlinked `site-packages` entries and entries that were copied (hardlinked), both modes used the same `pex-ns-pkgs.pth` strategy required for the symlink mode to support namespace packages at all. This strategy relies upon proper namespace packages though for distributions that use this feature. Unfortunately, this cannot be relied upon in practice and, as such, for a `--venv` with malformed namespace package dists, there is currently no recourse when attempting to create a `--venv` mode PEX. We can fix this by having the copy (hardlink) mode use the old merge strategy at the cost of more complex code.

The concrete example of malformed namespace package dists here were those produced by the Pants project for pantsbuild.pants and pantsbuild.pants.testutil between versions 2.5.0rc0 and 2.9.0rc1.


Discussion:

Resolving patch:
diff --git a/pex/tools/commands/venv.py b/pex/tools/commands/venv.py
--- a/pex/tools/commands/venv.py
+++ b/pex/tools/commands/venv.py
@@ -65,7 +65,8 @@ def _copytree(
         ):
             src_entry = os.path.join(root, path)
             dst_entry = os.path.join(dst, os.path.relpath(src_entry, src))
-            yield src_entry, dst_entry
+            if not is_dir:
+                yield src_entry, dst_entry
             try:
                 if symlink:
                     _relative_symlink(src_entry, dst_entry)
@@ -140,49 +141,45 @@ def record_provenance(src_to_dst):
     top_level_packages = Counter()  # type: typing.Counter[str]
     rel_extra_paths = OrderedSet()  # type: OrderedSet[str]
     for dist in pex.resolve():
-        # In the symlink case, in order to share all generated *.pyc files for a given distribution,
-        # we need to be able to have each contribution to a namespace package get its own top-level
-        # symlink. This requires adjoining extra sys.path entries beyond site-packages. We create
-        # the minimal number of extra such paths to satisfy all namespace package contributing dists
-        # for a given namespace package using a .pth file (See:
-        # https://docs.python.org/3/library/site.html).
-        #
-        # For example, given a PEX that depends on 3 different distributions contributing to the foo
-        # namespace package, we generate a layout like:
-        #   site-packages/
-        #     foo -> ../../../../../../installed_wheels/<hash>/foo-1.0-py3-none-any.why/foo
-        #     foo-1.0.dist-info -> ../../../../../../installed_wheels/<hash>/foo1/foo-1.0.dist-info
-        #     pex-ns-pkgs/
-        #       1/
-        #           foo -> ../../../../../../../../installed_wheels/<hash>/foo2-3.0-py3-none-any.whl/foo
-        #           foo2-3.0.dist-info -> ../../../../../../../../installed_wheels/<hash>/foo2-3.0-py3-none-any.whl/foo2-3.0.dist-info
-        #       2/
-        #           foo -> ../../../../../../../../installed_wheels/<hash>/foo3-2.5-py3-none-any.whl/foo
-        #           foo3-2.5.dist-info -> ../../../../../../../../installed_wheels/<hash>/foo3-2.5-py3-none-any.whl/foo2-2.5.dist-info
-        #     pex-ns-pkgs.pth
-        #
-        # Here site-packages/pex-ns-pkgs.pth contains:
-        #   pex-ns-pkgs/1
-        #   pex-ns-pkgs/2
-        #
-        # Although we don't need to do this for the link/copy case since directories (and
-        # subdirectories) then naturally merge, we do so anyhow to keep the code simpler and
-        # arguably provide a more transparently debuggable venv.
-        packages = [
-            name
-            for name in os.listdir(dist.location)
-            if name not in ("bin", "__pycache__")
-            and is_valid_python_identifier(name)
-            and os.path.isdir(os.path.join(dist.location, name))
-        ]
-        count = max(top_level_packages[package] for package in packages) if packages else 0
-        if count > 0:
-            rel_extra_path = os.path.join("pex-ns-pkgs", str(count))
-            dst = os.path.join(venv.site_packages_dir, rel_extra_path)
-            rel_extra_paths.add(rel_extra_path)
-        else:
-            dst = venv.site_packages_dir
-        top_level_packages.update(packages)
+        dst = venv.site_packages_dir
+        if symlink:
+            # In the symlink case, in order to share all generated *.pyc files for a given
+            # distribution, we need to be able to have each contribution to a namespace package get
+            # its own top-level symlink. This requires adjoining extra sys.path entries beyond
+            # site-packages. We create the minimal number of extra such paths to satisfy all
+            # namespace package contributing dists for a given namespace package using a .pth
+            # file (See: https://docs.python.org/3/library/site.html).
+            #
+            # For example, given a PEX that depends on 3 different distributions contributing to the
+            # foo namespace package, we generate a layout like:
+            #   site-packages/
+            #     foo -> ../../../../../../installed_wheels/<hash>/foo-1.0-py3-none-any.why/foo
+            #     foo-1.0.dist-info -> ../../../../../../installed_wheels/<hash>/foo1/foo-1.0.dist-info
+            #     pex-ns-pkgs/
+            #       1/
+            #           foo -> ../../../../../../../../installed_wheels/<hash>/foo2-3.0-py3-none-any.whl/foo
+            #           foo2-3.0.dist-info -> ../../../../../../../../installed_wheels/<hash>/foo2-3.0-py3-none-any.whl/foo2-3.0.dist-info
+            #       2/
+            #           foo -> ../../../../../../../../installed_wheels/<hash>/foo3-2.5-py3-none-any.whl/foo
+            #           foo3-2.5.dist-info -> ../../../../../../../../installed_wheels/<hash>/foo3-2.5-py3-none-any.whl/foo2-2.5.dist-info
+            #     pex-ns-pkgs.pth
+            #
+            # Here site-packages/pex-ns-pkgs.pth contains:
+            #   pex-ns-pkgs/1
+            #   pex-ns-pkgs/2
+            packages = [
+                name
+                for name in os.listdir(dist.location)
+                if name not in ("bin", "__pycache__")
+                and is_valid_python_identifier(name)
+                and os.path.isdir(os.path.join(dist.location, name))
+            ]
+            count = max(top_level_packages[package] for package in packages) if packages else 0
+            if count > 0:
+                rel_extra_path = os.path.join("pex-ns-pkgs", str(count))
+                dst = os.path.join(venv.site_packages_dir, rel_extra_path)
+                rel_extra_paths.add(rel_extra_path)
+            top_level_packages.update(packages)
 
         # N.B.: We do not include the top_level __pycache__ for a dist since there may be multiple
         # dists with top-level modules. In that case, one dists top-level __pycache__ would be


Test patch:
diff --git a/tests/integration/test_pex_bootstrapper.py b/tests/integration/test_pex_bootstrapper.py
--- a/tests/integration/test_pex_bootstrapper.py
+++ b/tests/integration/test_pex_bootstrapper.py
@@ -18,7 +18,7 @@
 from pex.typing import TYPE_CHECKING
 
 if TYPE_CHECKING:
-    from typing import Any, Callable
+    from typing import Any, Callable, Set, Text
 
 
 def test_ensure_venv_short_link(
@@ -122,7 +122,6 @@ def test_ensure_venv_namespace_packages(tmpdir):
     # type: (Any) -> None
 
     pex_root = os.path.join(str(tmpdir), "pex_root")
-    nspkgs_pex = os.path.join(str(tmpdir), "ns-pkgs.pex")
 
     # We know the twitter.common.metrics distributions depends on 4 other distributions contributing
     # to the twitter.common namespace package:
@@ -130,72 +129,93 @@ def test_ensure_venv_namespace_packages(tmpdir):
     # + twitter.common.decorators
     # + twitter.common.lang
     # + twitter.common.quantity
-    run_pex_command(
-        args=[
-            "twitter.common.metrics==0.3.11",
-            "-o",
-            nspkgs_pex,
-            "--runtime-pex-root",
-            pex_root,
-            "--venv",
-        ]
-    ).assert_success()
-    nspkgs_venv_pex = ensure_venv(PEX(nspkgs_pex), collisions_ok=False)
+    def create_ns_pkg_pex(copies):
+        # type: (bool) -> Virtualenv
+        nspkgs_pex = os.path.join(
+            str(tmpdir), "ns-pkgs-{style}.pex".format(style="copies" if copies else "symlinks")
+        )
+        run_pex_command(
+            args=[
+                "twitter.common.metrics==0.3.11",
+                "-o",
+                nspkgs_pex,
+                "--runtime-pex-root",
+                pex_root,
+                "--venv",
+                "--venv-site-packages-copies" if copies else "--no-venv-site-packages-copies",
+            ]
+        ).assert_success()
+        nspkgs_venv_pex = ensure_venv(PEX(nspkgs_pex), collisions_ok=False)
+
+        pex_info = PexInfo.from_pex(nspkgs_pex)
+        venv_dir = pex_info.venv_dir(nspkgs_pex)
+        assert venv_dir is not None
+        venv = Virtualenv(venv_dir=venv_dir)
+        assert os.path.realpath(nspkgs_venv_pex) == os.path.realpath(venv.join_path("pex"))
+        return venv
 
-    pex_info = PexInfo.from_pex(nspkgs_pex)
-    venv_dir = pex_info.venv_dir(nspkgs_pex)
-    assert venv_dir is not None
-    venv = Virtualenv(venv_dir=venv_dir)
+    venv_copies = create_ns_pkg_pex(copies=True)
+    assert not os.path.exists(os.path.join(venv_copies.site_packages_dir, "pex-ns-pkgs.pth"))
 
-    pex_ns_pkgs_pth = os.path.join(venv.site_packages_dir, "pex-ns-pkgs.pth")
+    venv_symlinks = create_ns_pkg_pex(copies=False)
+    pex_ns_pkgs_pth = os.path.join(venv_symlinks.site_packages_dir, "pex-ns-pkgs.pth")
     assert os.path.isfile(pex_ns_pkgs_pth)
     with open(pex_ns_pkgs_pth) as fp:
         assert 4 == len(fp.readlines())
 
     expected_path_entries = [
-        os.path.join(venv.site_packages_dir, d)
+        os.path.join(venv_symlinks.site_packages_dir, d)
         for d in ("", "pex-ns-pkgs/1", "pex-ns-pkgs/2", "pex-ns-pkgs/3", "pex-ns-pkgs/4")
     ]
     for d in expected_path_entries:
-        assert os.path.islink(os.path.join(venv.site_packages_dir, d, "twitter"))
-        assert os.path.isdir(os.path.join(venv.site_packages_dir, d, "twitter", "common"))
-
-    package_file_paths = set(
-        subprocess.check_output(
-            args=[
-                nspkgs_venv_pex,
-                "-c",
-                dedent(
-                    """\
-                    from __future__ import print_function
-                    import os
-
-                    from twitter.common import decorators, exceptions, lang, metrics, quantity
+        assert os.path.islink(os.path.join(venv_symlinks.site_packages_dir, d, "twitter"))
+        assert os.path.isdir(os.path.join(venv_symlinks.site_packages_dir, d, "twitter", "common"))
+
+    def find_package_paths(venv):
+        # type: (Virtualenv) -> Set[Text]
+        return set(
+            subprocess.check_output(
+                args=[
+                    venv.join_path("pex"),
+                    "-c",
+                    dedent(
+                        """\
+                        from __future__ import print_function
+                        import os
     
-                    
-                    for pkg in decorators, exceptions, lang, metrics, quantity:
-                        print(os.path.realpath(pkg.__file__))
-                    """
-                ),
-            ]
+                        from twitter.common import decorators, exceptions, lang, metrics, quantity
+        
+                        
+                        for pkg in decorators, exceptions, lang, metrics, quantity:
+                            # These are all packages; so __file__ looks like:
+                            #   <sys.path entry>/twitter/common/<pkg>/__init__.pyc
+                            print(os.path.realpath(os.path.dirname(os.path.dirname(pkg.__file__))))
+                        """
+                    ),
+                ]
+            )
+            .decode("utf-8")
+            .splitlines()
         )
-        .decode("utf-8")
-        .splitlines()
-    )
-    assert 5 == len(package_file_paths), "Expected 5 unique __init__.pyc files"
+
+    assert 1 == len(
+        find_package_paths(venv_copies)
+    ), "Expected 1 unique package path for a venv built from copies."
+
+    symlink_package_paths = find_package_paths(venv_symlinks)
+    assert 5 == len(symlink_package_paths), "Expected 5 unique package paths for symlinked venv."
 
     # We expect package paths like:
-    #   .../twitter.common.foo-0.3.11.*.whl/twitter/common/foo/__init__.pyc
+    #   .../twitter.common.foo-0.3.11.*.whl/twitter/common
     package_file_installed_wheel_dirs = {
-        os.path.dirname(os.path.dirname(os.path.dirname(os.path.dirname(p))))
-        for p in package_file_paths
+        os.path.dirname(os.path.dirname(p)) for p in symlink_package_paths
     }
-    assert os.path.realpath(pex_info.install_cache) == os.path.realpath(
+    assert os.path.realpath(os.path.join(pex_root, PexInfo.INSTALL_CACHE)) == os.path.realpath(
         os.path.commonprefix(list(package_file_installed_wheel_dirs))
     ), "Expected contributing wheel content to be symlinked from the installed wheel cache."
     assert {
         "twitter.common.{package}-0.3.11-py{py_major}-none-any.whl".format(
-            package=p, py_major=venv.interpreter.version[0]
+            package=p, py_major=venv_symlinks.interpreter.version[0]
         )
         for p in ("decorators", "exceptions", "lang", "metrics", "quantity")
     } == {