python_distribution editable installs in exports (#18639)

## About Editable Installs

Editable installs were traditionally done via pip with `pip install
--editable`. It is primarily useful during development when software
needs access to the entry points metadata.

When [PEP 517](https://peps.python.org/pep-0517/) was adopted, they
punted on how to allow for editable installs. [PEP
660](https://peps.python.org/pep-0660/) extended the PEP 517 backend API
to support building "editable" wheels. Therefore, there is now a
standard way to collect and install the metadata for "editable"
installs, using the "editable" wheel as an interchange between the
backend (which collects the metadata + builds the editable wheel) and
the frontend (which marshals the backend to perform a user-requested
"editable" install).

## Why would we need editable installs in pants?

I need editable installs in pants-exported virtualenvs so that dev tools
outside of pants have access to:
- The locked requirements
- The editable sources on the python path
- The entry points (and any other package metadata that can be loaded
from `dist-info`. Entry points is the biggest most impactful example)

I need to point all the external dev tooling at a virtualenv, and
technically I could export a pex that includes all of the
python-distributions pre-installed and use pex-tools to create a
virtualenv, but then I would have to recreate that venv for every dev
change wouldn't be a good dev experience.

One of those dev tools is `nosetest`. I considred using `run` to support
running that, but I am leary of adding all the complex BUILD machinery
to support running a tool that I'm trying to get rid of. Editable
installs is a more generic solution that serves my current needs, while
allowing for using it in other scenarios.

This PR comes in part from
#16621 (comment)

## Overview of this PR

### Scope & Future work

This PR focuses on adding editable installs to exported virtualenvs. Two
other issues ask for editable installs while running tests:
- #11386
- #15481

We can probably reuse a significant portion of this to generate editable
wheels for use in testing as well. Parts of this code will need to be
refactored to support that use case. But we also have to figure out the
UX for how users will define dependencies on a `python_distribution`
when they want an editable install instead of the built wheel to show up
in the sandbox. Anyway, addressing that is out of scope for this PR.

### New option `[export].py_editables_in_resolves` (a `StrListOption`)

This option allows user resolves to opt in to using the editable
installs. After
[consulting](https://pantsbuild.slack.com/archives/C0D7TNJHL/p1680810411706569?thread_ts=1680809713.310039&cid=C0D7TNJHL)
with @kaos, I decided to add an option for this instead of always trying
to generate/install the editable wheels.

> `python_distribution` does not have a `resolve` field. So figuring out
which resolve a `python_distribution` belongs to can be expensive:
calculating the owned deps of all distributions, and for each
distribution, look through those deps until one of them has a resolve
field, and use that for that dist’s resolve.
>
> Plus there’s the cost of building the PEP-660 wheels - if the
configured PEP-517 build backend doesn’t support the PEP-660 methods,
then it falls back to a method that is, sadly, optional in PEP-517. If
that method isn’t there, then it falls back to telling that backend to
build the whole wheel, and then it extracts just the dist-info directory
from it and discards the rest.
>
> So, installing these editable wheels isn’t free. It’ll slow down the
export, though I’m not sure by how much.

For StackStorm, I plan to set this in `pants.toml` for the default
resolve that has python_distributions.

Even without this option, I tried to bail out early if there were no
`python_distribution`s to install.

### Installing editable wheels for exports

I added this feature to the new export code path, which requires using
`export --resolve=`. The legacy codepath, which uses cli specs `export
<address specs>` did not change at all. I also ignored the `tool`
resolves which cannot have any relevant dists (and `tool` resolves are
deprecated anyway). Also, this is only for `mutable_virtualenv` exports,
as we need modify the virtualenv to install the editable wheels in the
venv after pex creates it from the lockfile.

When exporting a user resolve, we do a `Get(EditableLocalDists,
EditableLocalDistsRequest(resolve=resolve))`: _I'll skip over exactly
how this builds the wheels for now so this section can focus on how
installing works._


https://github.com/pantsbuild/pants/blob/f3a4620e81713f5022bf9a2dd1a4aa5ca100d1af/src/python/pants/backend/python/goals/export.py#L373-L379

As described in the commit message of
b5aa26a, I tried a variety of methods
for installing the editable wheels using pex. Ultimately, the best I
came up with is telling pex that the digest containing our editable
wheels are `sources` when building the `requirements.pex` used to
populate the venv, so that they would land in the virtualenv (even
though they land as plain wheel files.

Then we run `pex-tools` in a `PostProcessingCommand` to create and
populate the virtualenv, just as we did before this PR.

Once the virtualenv is created, we add 3 more `PostProcessingCommands`
to actually do the editable install. In this step, Pants is essentially
acting as the PEP-660 front end, though we use pip for some of the heavy
lifting. These commands:
1. move the editable wheels out of the virtualenv lib directory to the
temp dir that gets deleted at the end of the export
2. use pip to install all of the editable wheels (which contain a `.pth`
file that injects the source dir into `sys.path` and a `.dist-info`
directory with dist metadata such as entry points).
3. replace some of the pip-generated install metadata
(`*.dist-info/direct_url.json`) with our own so that we comply with
PEP-660 and mark the install as editable with a file url pointing the
the sources in the build_root (vs in a sandbox).

Now, anything that uses the exported venv should have access to the
standardized package metadata (in `.dist-info`) and the relevant source
roots should be automatically included in `sys.path`.

### Building PEP-660 editable wheels

The logic that actually builds the editable wheels is in
`pants.backend.python.util_rules.local_dists_pep660`. Building these
wheels requires the same chroot that pants uses to build regular wheels
and sdists. So, I refactored the rule in `util_rules.setup_py` so that I
could reuse the part that builds the `DistBuildRequest`.

These `local_dists_pep660` rules do approx this, starting with the rule
called in export:
- `Get(EditableLocalDists, EditableLocalDistsRequest(resolve=resolve))`
uses rule `build_editable_local_dists`
- injected arg: `ResolveSortedPythonDistributionTargets` comes from
rule: `sort_all_python_distributions_by_resolve`
- injected arg: `AllPythonDistributionTargets` comes from rule:
`find_all_python_distributions`
- `Get(LocalDistPEP660Wheels, PythonDistributionFieldSet.create(dist))`
for each dist in the resolve uses rule:
`isolate_local_dist_pep660_wheels`
- create `DistBuildRequest` using the `create_dist_build_request` method
I exposed in `util_rules.setup_py`
- `Get(PEP660BuildResult, DistBuildRequest)` uses rule:
`run_pep660_build`
            - generates the `.pth` file that goes in the editable wheel
            - runs a PEP 517 build backend wrapper script I wrote
- uses the PEP 517 build backend configured for the
`python_distribution` to generate the `.dist-info` directory
- generates the `WHEEL` and `RECORD` file to build a conformant wheel
file
- includes the `.pth` file previously generated (and placed in the
sandbox with the wrapper script)
- uses `zipfile` to build the wheel (using a vendored+modified function
from the `wheel` package).
                - prints a path to the generated wheel
- collects the generated editable wheel into a digest and collects
metadata about the digest similar to how the `local_dists` rules do.
- merges the editable wheel digests for all of the `python_distribution`
targets. This gets wrapped in `EditableLocalDists`

Much of the rule logic was based on (copied then modified):
`pants.backend.python.util_rules.dists` and
`pants.backend.python.util_rules.local_dists`.

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(),