python_distribution editable installs in exports

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

@@ -15,6 +15,10 @@
 from pants.backend.python.subsystems.setup import PythonSetup
 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 (
+    EditableLocalDists,
+    EditableLocalDistsRequest,
+)
 from pants.backend.python.util_rules.pex import Pex, PexProcess, PexRequest, VenvPex, VenvPexProcess
 from pants.backend.python.util_rules.pex_cli import PexPEX
 from pants.backend.python.util_rules.pex_environment import PexEnvironment
@@ -33,13 +37,13 @@
 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
 from pants.engine.target import Target
 from pants.engine.unions import UnionMembership, UnionRule, union
-from pants.option.option_types import BoolOption, EnumOption
+from pants.option.option_types import BoolOption, EnumOption, StrListOption
 from pants.util.docutil import bin_name
 from pants.util.strutil import path_safe, softwrap
 
@@ -123,6 +127,30 @@ class ExportPluginOptions:
         removal_hint="Set the `[export].py_resolve_format` option to 'symlinked_immutable_virtualenv'",
     )
 
+    py_editable_in_resolve = StrListOption(
+        # TODO: Is there a way to get [python].resolves in a memoized_property here?
+        #       If so, then we can validate that all resolves here are defined there.
+        help=softwrap(
+            """
+            When exporting a mutable virtualenv for a resolve, do PEP-660 editable installs
+            of all 'python_distribution' targets that own code in the exported resolve.
+
+            If a resolve name is not in this list, 'python_distribution' targets will not
+            be installed in the virtualenv. This defaults to an empty list for backwards
+            compatibility and to prevent unnecessary work to generate and install the
+            PEP-660 editable wheels.
+
+            This only applies when '[python].enable_resolves' is true and when exporting a
+            'mutable_virtualenv' ('symlinked_immutable_virtualenv' exports are not "full"
+            virtualenvs because they must not be edited, and do not include 'pip').
+
+            NOTE: If you are using legacy exports (not using the '--resolve' option), then
+            this option has no effect. Legacy exports will not include any editable installs.
+            """
+        ),
+        advanced=True,
+    )
+
 
 async def _get_full_python_version(pex_or_venv_pex: Pex | VenvPex) -> str:
     # Get the full python version (including patch #).
@@ -150,6 +178,7 @@ class VenvExportRequest:
     dest_prefix: str
     resolve_name: str
     qualify_path_with_python_version: bool
+    editable_local_dists_digest: Digest | None = None
 
 
 @rule
@@ -230,30 +259,92 @@ 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]),
+        ]
+
+        # Insert editable wheel post processing commands if needed.
+        if req.editable_local_dists_digest is not None:
+            # We need the snapshot to get the wheel file names which are something like:
+            #   - pkg_name-1.2.3-0.editable-py3-none-any.whl
+            wheels_snapshot = await Get(Snapshot, Digest, req.editable_local_dists_digest)
+            # We need the paths to the installed .dist-info directories to finish installation.
+            py_major_minor_version = ".".join(py_version.split(".", 2)[:2])
+            lib_dir = os.path.join(
+                output_path, "lib", f"python{py_major_minor_version}", "site-packages"
+            )
+            dist_info_dirs = [
+                # This builds: dist/.../resolve/3.8.9/lib/python3.8/site-packages/pkg_name-1.2.3.dist-info
+                os.path.join(lib_dir, "-".join(f.split("-")[:2]) + ".dist-info")
+                for f in wheels_snapshot.files
+            ]
+            # We use slice assignment to insert multiple elements at index 1.
+            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",
+                        "--no-deps",  # The deps were already installed via requirements.pex.
+                        "--no-build-isolation",  # Avoid VCS dep downloads (as they are installed).
+                        *(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 temp editable wheel)
+                        # with ours (which points to build_dir sources and is marked "editable").
+                        # Also update INSTALLER file to indicate that pants installed it.
+                        "sh",
+                        "-c",
+                        " ".join(
+                            [
+                                f"mv -f {src} {dst}; echo pants > {installer};"
+                                for src, dst, installer 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],
+                                    [os.path.join(d, "INSTALLER") 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:
@@ -269,8 +360,10 @@ class MaybeExportResult:
 async def export_virtualenv_for_resolve(
     request: _ExportVenvForResolveRequest,
     python_setup: PythonSetup,
+    export_subsys: ExportSubsystem,
     union_membership: UnionMembership,
 ) -> MaybeExportResult:
+    editable_local_dists_digest: Digest | None = None
     resolve = request.resolve
     lockfile_path = python_setup.resolves.get(resolve)
     if lockfile_path:
@@ -287,11 +380,20 @@ async def export_virtualenv_for_resolve(
             )
         )
 
+        if resolve in export_subsys.options.py_editable_in_resolve:
+            editable_local_dists = await Get(
+                EditableLocalDists, EditableLocalDistsRequest(resolve=resolve)
+            )
+            editable_local_dists_digest = editable_local_dists.optional_digest
+        else:
+            editable_local_dists_digest = None
+
         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,
@@ -337,6 +439,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-dist', version='1.2.3'),
+            dependencies=[':foo@resolve=a'],
+        )
+        """
+    ),
 }
 
 
@@ -97,7 +106,12 @@ def test_export(py_resolve_format: PythonResolveExportFormat) -> None:
     with setup_tmpdir(SOURCES) as tmpdir:
         resolve_names = ["a", "b", *(tool.name for tool in EXPORTED_TOOLS)]
         run_pants(
-            ["generate-lockfiles", "export", *(f"--resolve={name}" for name in resolve_names)],
+            [
+                "generate-lockfiles",
+                "export",
+                *(f"--resolve={name}" for name in resolve_names),
+                "--export-py-editable-in-resolve=['a']",
+            ],
             config=build_config(tmpdir, py_resolve_format),
         ).assert_success()
 
@@ -128,6 +142,29 @@ def test_export(py_resolve_format: PythonResolveExportFormat) -> None:
             expected_ansicolors_dir
         ), f"expected dist-info for ansicolors '{expected_ansicolors_dir}' does not exist"
 
+        if py_resolve_format == PythonResolveExportFormat.mutable_virtualenv:
+            expected_foo_dir = os.path.join(lib_dir, "foo_dist-1.2.3.dist-info")
+            if resolve == "b":
+                assert not os.path.isdir(
+                    expected_foo_dir
+                ), f"unexpected dist-info for foo-dist '{expected_foo_dir}' exists"
+            elif resolve == "a":
+                # make sure the editable wheel for the python_distribution is installed
+                assert os.path.isdir(
+                    expected_foo_dir
+                ), f"expected dist-info for foo-dist '{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-dist '{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-dist '{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

@@ -11,11 +11,17 @@
 from pants.backend.python.goals import export
 from pants.backend.python.goals.export import ExportVenvsRequest, PythonResolveExportFormat
 from pants.backend.python.lint.flake8 import subsystem as flake8_subsystem
-from pants.backend.python.target_types import PythonRequirementTarget
-from pants.backend.python.util_rules import pex_from_targets
+from pants.backend.python.macros.python_artifact import PythonArtifact
+from pants.backend.python.target_types import (
+    PythonDistribution,
+    PythonRequirementTarget,
+    PythonSourcesGeneratorTarget,
+)
+from pants.backend.python.util_rules import local_dists_pep660, pex_from_targets
 from pants.base.specs import RawSpecs, RecursiveGlobSpec
 from pants.core.goals.export import ExportResults
 from pants.core.util_rules import distdir
+from pants.engine.internals.parametrize import Parametrize
 from pants.engine.rules import QueryRule
 from pants.engine.target import Targets
 from pants.testutil.rule_runner import RuleRunner
@@ -30,11 +36,13 @@ def rule_runner() -> RuleRunner:
             *pex_from_targets.rules(),
             *target_types_rules.rules(),
             *distdir.rules(),
+            *local_dists_pep660.rules(),
             *flake8_subsystem.rules(),
             QueryRule(Targets, [RawSpecs]),
             QueryRule(ExportResults, [ExportVenvsRequest]),
         ],
-        target_types=[PythonRequirementTarget],
+        target_types=[PythonRequirementTarget, PythonSourcesGeneratorTarget, PythonDistribution],
+        objects={"python_artifact": PythonArtifact, "parametrize": Parametrize},
     )
 
 
@@ -170,8 +178,15 @@ def test_export_venv_new_codepath(
     current_interpreter = f"{vinfo.major}.{vinfo.minor}.{vinfo.micro}"
     rule_runner.write_files(
         {
+            "src/foo/__init__.py": "from colors import *",
             "src/foo/BUILD": dedent(
                 """\
+                python_sources(name='foo', resolve=parametrize('a', 'b'))
+                python_distribution(
+                    name='dist',
+                    provides=python_artifact(name='foo', version='1.2.3'),
+                    dependencies=[':foo@resolve=a'],
+                )
                 python_requirement(name='req1', requirements=['ansicolors==1.1.8'], resolve='a')
                 python_requirement(name='req2', requirements=['ansicolors==1.1.8'], resolve='b')
                 """
@@ -184,12 +199,16 @@ def test_export_venv_new_codepath(
     rule_runner.set_options(
         [
             f"--python-interpreter-constraints=['=={current_interpreter}']",
+            "--python-enable-resolves=True",
             "--python-resolves={'a': 'lock.txt', 'b': 'lock.txt'}",
             "--export-resolve=a",
             "--export-resolve=b",
             "--export-resolve=flake8",
             # Turn off lockfile validation to make the test simpler.
             "--python-invalid-lockfile-behavior=ignore",
+            # Turn off python synthetic lockfile targets to make the test simpler.
+            "--no-python-enable-lockfile-targets",
+            "--export-py-editable-in-resolve=['a', 'b']",
             format_flag,
         ],
         env_inherit={"PATH", "PYENV_ROOT"},
@@ -208,7 +227,13 @@ 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 == "a":
+                # editable wheels are installed for a user resolve that has dists
+                assert len(result.post_processing_cmds) == 5
+            else:
+                # tool resolves (flake8) and user resolves w/o dists (b)
+                # do not run the commands to do editable installs
+                assert len(result.post_processing_cmds) == 2
 
             ppc0 = result.post_processing_cmds[0]
             # The first arg is the full path to the python interpreter, which we
@@ -233,7 +258,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/register.py

@@ -46,6 +46,7 @@
 from pants.backend.python.util_rules import (
     ancestor_files,
     local_dists,
+    local_dists_pep660,
     pex,
     pex_from_targets,
     python_sources,
@@ -70,6 +71,7 @@ def rules():
         # Util rules
         *ancestor_files.rules(),
         *dependency_inference_rules.rules(),
+        *local_dists_pep660.rules(),
         *pex.rules(),
         *pex_from_targets.rules(),
         *python_sources.rules(),