Repository: pex-tool/pex

Issue:
Bad `installed_wheels` dir when a wheel's `top_level.txt` dirname collides with `sources`
Steps to reproduce:

Here's a `pex` (renamed to `.zip` for GitHub's sake) [local.zip](https://github.com/pantsbuild/pex/files/11805798/local.zip). It has a `top_level.txt` of `top_level`. Rename to `local.whl`

Then, next to that file:
```console
$ mkdir -p sources/top_level/
$ echo 'print("hello")' > sources/top_level/mymain.py
$ pex --output-file binary.pex --venv prepend --requirements-pex local.pex --entry-point top_level.mymain $'--sources-directory=sources' --layout packed
$ ./binary/__main__.py
hello
$ ls ~/.pex/installed_wheels/a571cccbc564d1e7da4efe2416aca4bb049256ab4d0b30553f9d4de8207fb283/top_level-0.0.0-cp38-cp38-linux_x86_64.whl/top_level/
__init__.py  mymain.py  __pycache__
```

`mymain.py` is not part of the wheel, and shouldn't be in there.






Discussion:
There is a bit of chaff in the example. The key bit is `--venv` and _not_ also `--venv-site-packages-copies` (for example, layout is irrelevant) - use `--venv-site-packages-copies`and all is well.

With just `--venv` you see:
```
$ ./binary.pex
hello
$ ls -l ~/.pex/venvs/s/fe373d35/venv/lib/python3.8/site-packages/
total 12
-rw-r--r-- 1 jsirois jsirois  230 Jun 30 18:14 PEX_EXTRA_SYS_PATH.pth
lrwxrwxrwx 1 jsirois jsirois  152 Jun 30 18:14 top_level -> ../../../../../../installed_wheels/a571cccbc564d1e7da4efe2416aca4bb049256ab4d0b30553f9d4de8207fb283/top_level-0.0.0-cp38-cp38-linux_x86_64.whl/top_level
drwxr-xr-x 2 jsirois jsirois 4096 Jun 30 18:14 top_level-0.0.0.dist-info
```

And that should be enough to see the issue. Now, for 3rdparty deps, this sort of collision (namespace package) is handled via a `pex-ns-pkgs.pth` you don't see here:
https://github.com/pantsbuild/pex/blob/c2fa8130d3dfb38f19f550a5bb0a3c69b7459f0d/pex/venv/installer.py#L459-L495
https://github.com/pantsbuild/pex/blob/c2fa8130d3dfb38f19f550a5bb0a3c69b7459f0d/pex/venv/installer.py#L514-L530

So, to handle this case where the namespace package is split across `--sources-directory` sources and 3rdaprty deps, the handling above will need to be de-localized.

The issue you're going to run into here though @thejcannon that Pex can't help you with is your example is not properly ns-packaged. The sources are (lack of `__init__.py`) but the local.zip does not cooperate - it has an empty `__init__.py`; so with the fix you'll get an error finding `top_level.mymain` (or not - vagaries of sys.path order dictate). I'll proceed with the Pex fix, but you'll have some work to do on your example or in Pants depending the source of the bad ns-package mix.

Resolving patch:
diff --git a/pex/cli/commands/venv.py b/pex/cli/commands/venv.py
--- a/pex/cli/commands/venv.py
+++ b/pex/cli/commands/venv.py
@@ -383,12 +383,14 @@ def _install_from_pex(
 
     if installer_configuration.scope in (InstallScope.ALL, InstallScope.DEPS_ONLY):
         if venv:
+            top_level_source_packages = tuple(installer.iter_top_level_source_packages(pex))
             installer.populate_venv_distributions(
                 venv=venv,
                 distributions=distributions,
                 provenance=provenance,
                 symlink=False,
                 hermetic_scripts=hermetic_scripts,
+                top_level_source_packages=top_level_source_packages,
             )
         else:
             installer.populate_flat_distributions(
diff --git a/pex/venv/installer.py b/pex/venv/installer.py
--- a/pex/venv/installer.py
+++ b/pex/venv/installer.py
@@ -31,7 +31,11 @@
 
 if TYPE_CHECKING:
     import typing
-    from typing import DefaultDict, Iterable, Iterator, List, Optional, Tuple, Union
+    from typing import Container, DefaultDict, Iterable, Iterator, List, Optional, Tuple, Union
+
+    import attr  # vendor:skip
+else:
+    from pex.third_party import attr
 
 
 def find_dist(
@@ -149,7 +153,7 @@ def _relative_symlink(
 def _copytree(
     src,  # type: str
     dst,  # type: str
-    exclude=(),  # type: Tuple[str, ...]
+    exclude=(),  # type: Container[str]
     symlink=False,  # type: bool
 ):
     # type: (...) -> Iterator[Tuple[str, str]]
@@ -333,6 +337,7 @@ def populate_venv_distributions(
     provenance,  # type: Provenance
     symlink=False,  # type: bool
     hermetic_scripts=True,  # type: bool
+    top_level_source_packages=(),  # type: Iterable[str]
 ):
     # type: (...) -> None
 
@@ -343,6 +348,7 @@ def populate_venv_distributions(
             venv_python=provenance.target_python,
             symlink=symlink,
             hermetic_scripts=hermetic_scripts,
+            top_level_source_packages=top_level_source_packages,
         )
     )
 
@@ -378,6 +384,27 @@ def populate_venv_sources(
     return shebang
 
 
+def _iter_top_level_packages(
+    path,  # type: str
+    excludes=("bin", "__pycache__"),  # type: Container[str]
+):
+    # type: (...) -> Iterator[str]
+    for name in os.listdir(path):
+        if (
+            name not in excludes
+            and is_valid_python_identifier(name)
+            and os.path.isdir(os.path.join(path, name))
+        ):
+            yield name
+
+
+def iter_top_level_source_packages(pex):
+    # type: (PEX) -> Iterator[str]
+    pex_sources = PEXSources.mount(pex)
+    for path in _iter_top_level_packages(pex_sources.path, excludes=pex_sources.excludes):
+        yield path
+
+
 def populate_venv_from_pex(
     venv,  # type: Virtualenv
     pex,  # type: PEX
@@ -392,6 +419,7 @@ def populate_venv_from_pex(
 
     provenance = Provenance.create(venv, python=python)
     shebang = provenance.calculate_shebang(hermetic_scripts=hermetic_scripts)
+    top_level_source_packages = tuple(iter_top_level_source_packages(pex))
 
     if scope in (InstallScope.ALL, InstallScope.DEPS_ONLY):
         populate_venv_distributions(
@@ -400,6 +428,7 @@ def populate_venv_from_pex(
             symlink=symlink,
             hermetic_scripts=hermetic_scripts,
             provenance=provenance,
+            top_level_source_packages=top_level_source_packages,
         )
 
     if scope in (InstallScope.ALL, InstallScope.SOURCE_ONLY):
@@ -446,13 +475,14 @@ def _populate_venv_deps(
     venv_python,  # type: str
     symlink=False,  # type: bool
     hermetic_scripts=True,  # type: bool
+    top_level_source_packages=(),  # type: Iterable[str]
 ):
     # type: (...) -> Iterator[Tuple[str, str]]
 
     # Since the pex distributions are all materialized to ~/.pex/installed_wheels, which we control,
     # we can optionally symlink to take advantage of sharing generated *.pyc files for auto-venvs
     # created in ~/.pex/venvs.
-    top_level_packages = Counter()  # type: typing.Counter[str]
+    top_level_packages = Counter(top_level_source_packages)  # type: typing.Counter[str]
     rel_extra_paths = OrderedSet()  # type: OrderedSet[str]
     for dist in distributions:
         rel_extra_path = None
@@ -481,13 +511,7 @@ def _populate_venv_deps(
             # 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))
-            ]
+            packages = list(_iter_top_level_packages(dist.location))
             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))
@@ -536,6 +560,28 @@ def _populate_venv_deps(
         TRACER.log("Re-writing {}".format(script))
 
 
+@attr.s(frozen=True)
+class PEXSources(object):
+    @classmethod
+    def mount(cls, pex):
+        # type: (PEX) -> PEXSources
+        return cls(
+            path=PEXEnvironment.mount(pex.path()).path,
+            excludes=(
+                "__main__.py",
+                "__pex__",
+                "__pycache__",
+                layout.BOOTSTRAP_DIR,
+                layout.DEPS_DIR,
+                layout.PEX_INFO_PATH,
+                layout.PEX_LAYOUT_PATH,
+            ),
+        )
+
+    path = attr.ib()  # type: str
+    excludes = attr.ib()  # type: Container[str]
+
+
 def _populate_sources(
     pex,  # type: PEX
     dst,  # type: str
@@ -544,18 +590,11 @@ def _populate_sources(
 
     # Since the pex.path() is ~always outside our control (outside ~/.pex), we copy all PEX user
     # sources into the venv.
+    pex_sources = PEXSources.mount(pex)
     for src, dst in _copytree(
-        src=PEXEnvironment.mount(pex.path()).path,
+        src=pex_sources.path,
         dst=dst,
-        exclude=(
-            "__main__.py",
-            "__pex__",
-            "__pycache__",
-            layout.BOOTSTRAP_DIR,
-            layout.DEPS_DIR,
-            layout.PEX_INFO_PATH,
-            layout.PEX_LAYOUT_PATH,
-        ),
+        exclude=pex_sources.excludes,
         symlink=False,
     ):
         yield src, dst
diff --git a/pex/venv/virtualenv.py b/pex/venv/virtualenv.py
--- a/pex/venv/virtualenv.py
+++ b/pex/venv/virtualenv.py
@@ -349,8 +349,7 @@ def sys_path(self):
 
     def iter_distributions(self):
         # type: () -> Iterator[Distribution]
-        """"""
-        for dist in find_distributions(search_path=self.sys_path):
+        for dist in find_distributions(search_path=self._interpreter.site_packages):
             yield dist
 
     def _rewrite_base_scripts(self, real_venv_dir):


Test patch:
diff --git a/tests/integration/cli/commands/test_issue_1667.py b/tests/integration/cli/commands/test_issue_1667.py
--- a/tests/integration/cli/commands/test_issue_1667.py
+++ b/tests/integration/cli/commands/test_issue_1667.py
@@ -77,14 +77,19 @@ def test_interpreter_constraints_range_coverage(
 
     def assert_pex_works(python):
         # type: (str) -> None
-        assert (
-            subprocess.check_output(
-                args=[python, ipython_pex, "-c", "import IPython; print(IPython.__file__)"]
-            )
-            .decode("utf-8")
-            .strip()
-            .startswith(pex_root)
+        comm_channel = os.path.join(str(tmpdir), "comm")
+        subprocess.check_call(
+            args=[
+                python,
+                ipython_pex,
+                "-c",
+                "import IPython; print(IPython.__file__, file=open({comm!r}, 'w'))".format(
+                    comm=comm_channel
+                ),
+            ]
         )
+        with open(comm_channel) as fp:
+            assert fp.read().startswith(pex_root)
 
     assert_pex_works(py38.binary)
     if (3, 7) <= sys.version_info[:2] < (3, 11):
diff --git a/tests/integration/test_issue_1726.py b/tests/integration/test_issue_1726.py
--- a/tests/integration/test_issue_1726.py
+++ b/tests/integration/test_issue_1726.py
@@ -46,6 +46,13 @@ def test_check_install_issue_1726(
             )
         )
 
+    # Via: jaraco-collections==3.5.1 -> jaraco-text -> inflect -> pydantic>=1.9.1
+    # Pydantic 2.0 can get pulled in which defeats the Pip legacy resolver and leads to a resolve
+    # conflict for PyPy; so we bound pydantic low.
+    constraints = os.path.join(str(tmpdir), "constraints.txt")
+    with open(constraints, "w") as fp:
+        fp.write("pydantic<2")
+
     pex_root = os.path.join(str(tmpdir), "pex_root")
     pex_args = [
         "--pex-root",
@@ -53,6 +60,8 @@ def test_check_install_issue_1726(
         "--runtime-pex-root",
         pex_root,
         src,
+        "--constraints",
+        constraints,
         "--",
         "-c",
         "from jaraco import collections; print(collections.__file__)",
diff --git a/tests/integration/venv_ITs/test_issue_2160.py b/tests/integration/venv_ITs/test_issue_2160.py
new file mode 100644
--- /dev/null
+++ b/tests/integration/venv_ITs/test_issue_2160.py
@@ -0,0 +1,141 @@
+# Copyright 2023 Pants project contributors (see CONTRIBUTORS.md).
+# Licensed under the Apache License, Version 2.0 (see LICENSE).
+import glob
+import os.path
+import shutil
+import subprocess
+import sys
+from textwrap import dedent
+
+import pytest
+
+from pex.common import safe_open, safe_rmtree
+from pex.layout import Layout
+from pex.pep_440 import Version
+from pex.pep_503 import ProjectName
+from pex.pex import PEX
+from pex.testing import WheelBuilder, run_pex_command
+from pex.typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from typing import Any, List
+
+
+@pytest.fixture
+def top_level_wheel(tmpdir):
+    # type: (Any) -> str
+    top_level_project = os.path.join(str(tmpdir), "project")
+    with safe_open(os.path.join(top_level_project, "top_level", "__init__.py"), "w") as fp:
+        fp.write("__path__ = __import__('pkgutil').extend_path(__path__, __name__)")
+    with safe_open(os.path.join(top_level_project, "top_level", "lib.py"), "w") as fp:
+        fp.write("OG = 'Henry Barber'")
+
+    with safe_open(os.path.join(top_level_project, "setup.cfg"), "w") as fp:
+        fp.write(
+            dedent(
+                """\
+                [metadata]
+                name = top_level
+                version = 0.1.0
+
+                [options]
+                packages =
+                    top_level
+                """
+            )
+        )
+    with safe_open(os.path.join(top_level_project, "pyproject.toml"), "w") as fp:
+        fp.write(
+            dedent(
+                """\
+                [build-system]
+                requires = ["setuptools"]
+                build-backend = "setuptools.build_meta"
+                """
+            )
+        )
+
+    wheel_dir = os.path.join(str(tmpdir), "wheels")
+    return WheelBuilder(source_dir=top_level_project, wheel_dir=wheel_dir).bdist()
+
+
+@pytest.mark.parametrize(
+    "style_args",
+    [
+        pytest.param([], id="zipapp"),
+        pytest.param(["--venv"], id="venv (site-packages symlinks)"),
+        pytest.param(["--venv", "--venv-site-packages-copies"], id="venv (site-packages copies)"),
+    ],
+)
+@pytest.mark.parametrize(
+    "layout", [pytest.param(layout, id=layout.value) for layout in Layout.values()]
+)
+def test_ns_package_split_across_sources_and_deps(
+    tmpdir,  # type: Any
+    top_level_wheel,  # type: str
+    style_args,  # type: List[str]
+    layout,  # type: Layout.Value
+):
+    # type: (...) -> None
+
+    sources = os.path.join(str(tmpdir), "sources")
+    with safe_open(os.path.join(sources, "top_level", "__init__.py"), "w") as fp:
+        fp.write("__path__ = __import__('pkgutil').extend_path(__path__, __name__)")
+    with safe_open(os.path.join(sources, "top_level", "mymain.py"), "w") as fp:
+        fp.write(
+            dedent(
+                """\
+                from top_level.lib import OG
+
+
+                print("Hello {}!".format(OG))
+                """
+            )
+        )
+
+    pex_root = os.path.join(str(tmpdir), "pex_root")
+    pex = os.path.join(str(tmpdir), "binary.pex")
+    run_pex_command(
+        args=[
+            "--pex-root",
+            pex_root,
+            "--runtime-pex-root",
+            pex_root,
+            top_level_wheel,
+            "-D",
+            sources,
+            "-e",
+            "top_level.mymain",
+            "-o",
+            pex,
+            "--layout",
+            layout.value,
+        ]
+        + style_args
+    ).assert_success()
+
+    safe_rmtree(pex_root)
+    assert (
+        "Hello Henry Barber!"
+        == subprocess.check_output(args=[sys.executable, pex]).decode("utf-8").strip()
+    )
+
+    distributions = list(PEX(pex).resolve())
+    assert 1 == len(distributions)
+
+    top_level = distributions[0]
+    assert ProjectName("top_level") == top_level.metadata.project_name
+    assert Version("0.1.0") == top_level.metadata.version
+
+    assert sorted(
+        os.path.join(top_level.location, "top_level", expected)
+        for expected in ("__init__.py", "lib.py")
+    ) == sorted(
+        os.path.join(root, f)
+        for root, dirs, files in os.walk(top_level.location)
+        for f in files
+        if f.endswith(".py")
+    ), (
+        "Even in venv symlink mode, we expect the PEX 3rd-party wheel sources to remain isolated "
+        "from the PEX user sources"
+    )