refactor editable wheel passing and installation

Installing the editable wheels via a pex was not effective.
Using pex-tools to install two different pexes in the same venv
caused conflicts with the __main__.py file. Using pex_path did
not work because the editable wheel pex only had the editable
wheels, not any of the dependencies. So, installing them in
a pex was not effective.

However, it does work to tell pex that the wheels are "sources".
That way, pex just dumps them in the lib directory. Then, it
was just a matter of adding the post-processing commands to
move, install, and wrap up installation of the editable wheels.

Even though we use pip to install the editable wheels, pants is
both backend and frontend for PEP 660, because we generate the
wheel (delegating dist-info generation to the external backend)
and then install the wheel (delegating unpacking the wheel to
pip). The final piece of complying with PEP-660 is to mark the
install as an editable install via direct_url.json. We generate
that when building the wheel, and then move it into place after
pip has installed the editable wheels for us.

One other method I considered was installing the editable
wheels in the requirements pex. But I could not figure out how
to combine EntireLockfile with additional wheel requirements.
So, I punted and went merely passing the wheels as is via pex
instead of attempting to install them in any pex.

Finally, all of the export and local_dists_pep660 test pass!

src/python/pants/backend/python/goals/export.py

@@ -16,8 +16,8 @@
 from pants.backend.python.target_types import PexLayout, PythonResolveField
 from pants.backend.python.util_rules.interpreter_constraints import InterpreterConstraints
 from pants.backend.python.util_rules.local_dists_pep660 import (
-    LocalDistsPEP660Pex,
-    LocalDistsPEP660PexRequest,
+    EditableLocalDists,
+    EditableLocalDistsRequest,
 )
 from pants.backend.python.util_rules.local_dists_pep660 import rules as local_dists_pep660_rules
 from pants.backend.python.util_rules.pex import Pex, PexProcess, PexRequest, VenvPex, VenvPexProcess
@@ -38,7 +38,7 @@
 from pants.core.util_rules.distdir import DistDir
 from pants.engine.engine_aware import EngineAwareParameter
 from pants.engine.environment import EnvironmentName
-from pants.engine.internals.native_engine import AddPrefix, Digest, MergeDigests
+from pants.engine.internals.native_engine import AddPrefix, Digest, MergeDigests, Snapshot
 from pants.engine.internals.selectors import Get, MultiGet
 from pants.engine.process import ProcessCacheScope, ProcessResult
 from pants.engine.rules import collect_rules, rule
@@ -155,6 +155,7 @@ class VenvExportRequest:
     dest_prefix: str
     resolve_name: str
     qualify_path_with_python_version: bool
+    editable_local_dists_digest: Digest | None = None
 
 
 @rule
@@ -235,30 +236,81 @@ async def do_export(
         tmpdir_under_digest_root = os.path.join("{digest_root}", tmpdir_prefix)
         merged_digest_under_tmpdir = await Get(Digest, AddPrefix(merged_digest, tmpdir_prefix))
 
+        post_processing_cmds = [
+            PostProcessingCommand(
+                complete_pex_env.create_argv(
+                    os.path.join(tmpdir_under_digest_root, pex_pex.exe),
+                    *[
+                        os.path.join(tmpdir_under_digest_root, requirements_pex.name),
+                        "venv",
+                        "--pip",
+                        "--collisions-ok",
+                        output_path,
+                    ],
+                ),
+                {
+                    **complete_pex_env.environment_dict(python=requirements_pex.python),
+                    "PEX_MODULE": "pex.tools",
+                },
+            ),
+            # Remove the requirements and pex pexes, to avoid confusion.
+            PostProcessingCommand(["rm", "-rf", tmpdir_under_digest_root]),
+        ]
+
+        if req.editable_local_dists_digest:
+            # insert editable wheel post processing commands
+            wheels_snapshot = await Get(Snapshot, Digest, req.editable_local_dists_digest)
+            py_minor_version = ".".join(py_version.split(".")[:2])
+            lib_dir = os.path.join(output_path, "lib", f"python{py_minor_version}", "site-packages")
+            dist_info_dirs = [
+                os.path.join(lib_dir, "-".join(f.split("-")[:2]) + ".dist-info")
+                for f in wheels_snapshot.files
+            ]
+            post_processing_cmds[1:1] = [
+                PostProcessingCommand(
+                    [
+                        # the wheels are "sources" in the pex and get dumped in lib_dir
+                        # so we move them to tmpdir where they will be removed at the end.
+                        "mv",
+                        *[os.path.join(lib_dir, f) for f in wheels_snapshot.files],
+                        tmpdir_under_digest_root,
+                    ]
+                ),
+                PostProcessingCommand(
+                    [
+                        # now install the editable wheels
+                        os.path.join(output_path, "bin", "pip"),
+                        "install",
+                        *[os.path.join(tmpdir_under_digest_root, f) for f in wheels_snapshot.files],
+                    ]
+                ),
+                PostProcessingCommand(
+                    [
+                        # replace pip's direct_url.json (which points to the wheel)
+                        # with ours (which points to build_dir sources and is marked "editable").
+                        "sh",
+                        "-c",
+                        " ".join(
+                            [
+                                f"mv -f {src} {dst};"
+                                for src, dst in zip(
+                                    [
+                                        os.path.join(d, "direct_url__pants__.json")
+                                        for d in dist_info_dirs
+                                    ],
+                                    [os.path.join(d, "direct_url.json") for d in dist_info_dirs],
+                                )
+                            ]
+                        ),
+                    ]
+                ),
+            ]
+
         return ExportResult(
             description,
             dest,
             digest=merged_digest_under_tmpdir,
-            post_processing_cmds=[
-                PostProcessingCommand(
-                    complete_pex_env.create_argv(
-                        os.path.join(tmpdir_under_digest_root, pex_pex.exe),
-                        *[
-                            os.path.join(tmpdir_under_digest_root, requirements_pex.name),
-                            "venv",
-                            "--pip",
-                            "--collisions-ok",
-                            output_path,
-                        ],
-                    ),
-                    {
-                        **complete_pex_env.environment_dict(python=requirements_pex.python),
-                        "PEX_MODULE": "pex.tools",
-                    },
-                ),
-                # Remove the requirements and pex pexes, to avoid confusion.
-                PostProcessingCommand(["rm", "-rf", tmpdir_under_digest_root]),
-            ],
+            post_processing_cmds=post_processing_cmds,
             resolve=req.resolve_name or None,
         )
     else:
@@ -276,6 +328,7 @@ async def export_virtualenv_for_resolve(
     python_setup: PythonSetup,
     union_membership: UnionMembership,
 ) -> MaybeExportResult:
+    editable_local_dists_digest: Digest | None = None
     resolve = request.resolve
     lockfile_path = python_setup.resolves.get(resolve)
     if lockfile_path:
@@ -292,27 +345,20 @@ async def export_virtualenv_for_resolve(
             )
         )
 
-        pex_path = []
-
         editable_local_dists = await Get(
-            LocalDistsPEP660Pex,
-            LocalDistsPEP660PexRequest(
-                interpreter_constraints=interpreter_constraints,
-                resolve=resolve,
-            ),
+            EditableLocalDists, EditableLocalDistsRequest(resolve=resolve)
         )
-        if editable_local_dists.pex:
-            pex_path.append(editable_local_dists.pex)
+        editable_local_dists_digest = editable_local_dists.optional_digest
 
         pex_request = PexRequest(
             description=f"Build pex for resolve `{resolve}`",
             output_filename=f"{path_safe(resolve)}.pex",
             internal_only=True,
             requirements=EntireLockfile(lockfile),
+            sources=editable_local_dists_digest,
             interpreter_constraints=interpreter_constraints,
             # Packed layout should lead to the best performance in this use case.
             layout=PexLayout.PACKED,
-            pex_path=pex_path,
         )
     else:
         # It's a tool resolve.
@@ -355,6 +401,7 @@ async def export_virtualenv_for_resolve(
             dest_prefix,
             resolve,
             qualify_path_with_python_version=True,
+            editable_local_dists_digest=editable_local_dists_digest,
         ),
     )
     return MaybeExportResult(export_result)

src/python/pants/backend/python/goals/export_integration_test.py

@@ -24,7 +24,16 @@
         """
     ),
     "src/python/foo.py": "from colors import *",
-    "src/python/BUILD": "python_source(name='foo', source='foo.py', resolve=parametrize('a', 'b'))",
+    "src/python/BUILD": dedent(
+        """\
+        python_source(name='foo', source='foo.py', resolve=parametrize('a', 'b'))
+        python_distribution(
+            name='dist',
+            provides=python_artifact(name='foo', version='1.2.3'),
+            dependencies=[':foo@resolve=a'],
+        )
+        """
+    ),
 }
 
 
@@ -128,6 +137,24 @@ def test_export(py_resolve_format: PythonResolveExportFormat) -> None:
             expected_ansicolors_dir
         ), f"expected dist-info for ansicolors '{expected_ansicolors_dir}' does not exist"
 
+        if resolve == "a" and py_resolve_format == PythonResolveExportFormat.mutable_virtualenv:
+            # make sure the editable wheel for the python_distribution is installed
+            expected_foo_dir = os.path.join(lib_dir, "foo-1.2.3.dist-info")
+            assert os.path.isdir(
+                expected_foo_dir
+            ), f"expected dist-info for foo '{expected_foo_dir}' does not exist"
+            # direct_url__pants__.json should be moved to direct_url.json
+            expected_foo_direct_url_pants = os.path.join(
+                expected_foo_dir, "direct_url__pants__.json"
+            )
+            assert not os.path.isfile(
+                expected_foo_direct_url_pants
+            ), f"expected direct_url__pants__.json for foo '{expected_foo_direct_url_pants}' was not removed"
+            expected_foo_direct_url = os.path.join(expected_foo_dir, "direct_url.json")
+            assert os.path.isfile(
+                expected_foo_direct_url
+            ), f"expected direct_url.json for foo '{expected_foo_direct_url}' does not exist"
+
     for tool_config in EXPORTED_TOOLS:
         export_dir = os.path.join(export_prefix, tool_config.name, platform.python_version())
         assert os.path.isdir(export_dir), f"expected export dir '{export_dir}' does not exist"

src/python/pants/backend/python/goals/export_test.py

@@ -208,7 +208,11 @@ def test_export_venv_new_codepath(
             assert ppc1.argv[3] == "{digest_root}"
             assert ppc1.extra_env == FrozenDict()
         else:
-            assert len(result.post_processing_cmds) == 2
+            if resolve == "flake8":
+                assert len(result.post_processing_cmds) == 2
+            else:
+                # editable wheels are installed for a user resolve
+                assert len(result.post_processing_cmds) == 5
 
             ppc0 = result.post_processing_cmds[0]
             # The first arg is the full path to the python interpreter, which we
@@ -233,7 +237,7 @@ def test_export_venv_new_codepath(
             assert ppc0.extra_env["PEX_MODULE"] == "pex.tools"
             assert ppc0.extra_env.get("PEX_ROOT") is not None
 
-            ppc1 = result.post_processing_cmds[1]
+            ppc1 = result.post_processing_cmds[-1]
             assert ppc1.argv == ("rm", "-rf", tmpdir)
             assert ppc1.extra_env == FrozenDict()
 

src/python/pants/backend/python/util_rules/local_dists_pep660.py

@@ -26,9 +26,7 @@
     distutils_repr,
 )
 from pants.backend.python.util_rules.dists import rules as dists_rules
-from pants.backend.python.util_rules.interpreter_constraints import InterpreterConstraints
-from pants.backend.python.util_rules.pex import Pex, PexRequest, VenvPex, VenvPexProcess
-from pants.backend.python.util_rules.pex_requirements import PexRequirements
+from pants.backend.python.util_rules.pex import PexRequest, VenvPex, VenvPexProcess
 from pants.base.build_root import BuildRoot
 from pants.core.util_rules import system_binaries
 from pants.core.util_rules.system_binaries import BashBinary, UnzipBinary
@@ -463,58 +461,47 @@ async def sort_all_python_distributions_by_resolve(
 
 
 @dataclass(frozen=True)
-class LocalDistsPEP660PexRequest:  # based on LocalDistsPexRequest
-    """Request to build a PEX populated by PEP660 wheels of local dists.
-
-    Like LocalDistsPexRequest, the local dists come from the dependency closure of a set of
-    addresses. Unlike LocalDistsPexRequest, the editable wheel files must not be exported or made
-    available to the end-user (according to PEP 660). Instead, the PEP660Pex serves as an
-    intermediate, internal-only, PEX that can be used to install these wheels in a virtualenv. It
-    follows then, that this PEX should probably not be exported for use by end-users.
+class EditableLocalDistsRequest:
+    """Request toild generate PEP660 wheels of local dists in the given resolve.
+
+    The editable wheel files must not be exported or made available to the end-user (according to
+    PEP 660). Instead, the PEP660 editable wheels serve as intermediate, internal-only,
+    representation of what should be installed in the exported virtualenv to create the editable
+    installs of local python_distributions.
     """
 
-    interpreter_constraints: InterpreterConstraints
     resolve: str | None  # None if resolves is not enabled
 
 
 @dataclass(frozen=True)
-class LocalDistsPEP660Pex:  # based on LocalDistsPex
-    """A PEX file populated by PEP660 wheels of local dists.
+class EditableLocalDists:
+    """A Digest populated by editable (PEP660) wheels of local dists.
 
-    Can be consumed from another PEX, e.g., by adding to PEX_PATH.
+    According to PEP660, these wheels should not be exported to users and must be discarded
+    after install. Anything that uses this should ensure that these wheels get installed and
+    then deleted.
 
-    The PEX will contain installs of PEP660 wheels generate4d from local dists. Installing PEP660
-    wheels creates an "editable" install such that the sys.path gets adjusted to include source
-    directories that are NOT part of the PEX. As such, this PEX is decidedly not hermetic or
-    portable and should only be used to facilitate building virtualenvs for development.
+    Installing PEP660 wheels creates an "editable" install such that the sys.path gets
+    adjusted to include source directories from the build root (not from the sandbox).
+    This is decidedly not hermetic or portable and should only be used locally.
 
     PEP660 wheels have .dist-info metadata and the .pth files (or similar) that adjust sys.path.
-
-    The PEX will only contain metadata for local dists and not any dependencies. For Pants generated
-    `setup.py` / `pyproject.toml`, the dependencies will be included in the standard resolve process
-    that the locally-built dists PEX is adjoined to via PEX_PATH. For hand-made `setup.py` /
-    `pyproject.toml` with 3rdparty dependencies not hand-mirrored into BUILD file dependencies, this
-    will lead to issues. See https://github.com/pantsbuild/pants/issues/13587#issuecomment-974863636
-    for one way to fix this corner which is intentionally punted on for now.
-
-    Lists the files provided by the dists on sys.path, so they can be subtracted from
-    sources digests, to prevent the same file ending up on sys.path twice.
     """
 
-    pex: Pex | None
+    optional_digest: Digest | None
 
 
 @rule(desc="Building editable local distributions (PEP 660)")
 async def build_editable_local_dists(  # based on build_local_dists
-    request: LocalDistsPEP660PexRequest,
+    request: EditableLocalDistsRequest,
     all_dists: ResolveSortedPythonDistributionTargets,
     python_setup: PythonSetup,
-) -> LocalDistsPEP660Pex:
+) -> EditableLocalDists:
     resolve = request.resolve if python_setup.enable_resolves else None
     resolve_dists = all_dists.targets.get(resolve, ())
 
     if not resolve_dists:
-        return LocalDistsPEP660Pex(None)
+        return EditableLocalDists(None)
 
     local_dists_wheels = await MultiGet(
         Get(
@@ -533,19 +520,7 @@ async def build_editable_local_dists(  # based on build_local_dists
 
     wheels_digest = await Get(Digest, MergeDigests(wheels_digests))
 
-    editable_dists_pex = await Get(
-        Pex,
-        PexRequest(
-            output_filename="editable_local_dists.pex",
-            requirements=PexRequirements(wheels),
-            interpreter_constraints=request.interpreter_constraints,
-            additional_inputs=wheels_digest,
-            internal_only=True,
-            additional_args=["--intransitive"],
-        ),
-    )
-
-    return LocalDistsPEP660Pex(editable_dists_pex)
+    return EditableLocalDists(wheels_digest)
 
 
 def rules():