A --resolve flag, to specify the resolves to export. (Cherry-pick of #…

…17416) (#17461)

Currently, export takes cmd-line specs, so the python export implementation, which
exports virtualenvs, exports the requirements of those targets. 

This has a couple of issues:

A) User intent is almost certainly that export should emit the entire lockfile, whereas we will omit anything
  that isn't actively depended on by some target in the specs (which could be the case even with `::`).
B) Tools don't relate to specs, so we export every tool every time, no matter the specs, which is obviously clunky.

This change re-envisions how export should work. Instead of passing specs, you pass the `--resolve` flag one or more times, and we export a venv for each resolve, whether it's a user resolve or a tool resolve. 

The old codepath still works, but we can consider deprecating it in a followup.

Closes #17398

docs/markdown/Using Pants/setting-up-an-ide.md

@@ -34,29 +34,25 @@ See [Use of the PYTHONPATH variable](https://code.visualstudio.com/docs/python/e
 Python third-party dependencies and tools
 -----------------------------------------
 
-To get your editor to understand the repo's third-party dependencies, you will probably want to point it at a virtualenv containing those dependencies.
+To get your editor to understand the repo's third-party Python dependencies, you will probably want to point it at a virtualenv containing those dependencies.
 
-You can use the `export` goal to create a suitable virtualenv. 
+Assuming you are using the ["resolves" feature for Python lockfiles](doc:python-third-party-dependencies)—which we strongly recommend—Pants can export a virtualenv for each of your resolves. You can then point your IDE to whichever resolve you want to load at the time.
+
+To use the `export` goal to create a virtualenv:
 
 ```
-❯ ./pants export ::
-Wrote virtualenv for the resolve 'python-default' (using CPython==3.9.*) to dist/export/python/virtualenvs/python-default
+❯ ./pants export --symlink-python-virtualenv --resolve=python-default
+Wrote symlink to immutable virtualenv for python-default (using Python 3.9.13) to dist/export/python/virtualenvs/python-default
 ```
 
-If you are using the ["resolves" feature for Python lockfiles](doc:python-third-party-dependencies)—which we strongly recommend—Pants will write the virtualenv to `dist/export/python/virtualenvs/<resolve-name>`. If you have multiple resolves, this means that Pants will create one virtualenv per resolve. You can then point your IDE to whichever resolve you want to load at the time.
+You can specify the `--resolve` flag [multiple times](doc:options#list-values) to export multiple virtualenvs at once.
 
-### Tool virtualenvs
+The `--symlink-python-virtualenv` option symlinks to an immutable, internal virtualenv that does not have `pip` installed in it. This method is faster, but you must be careful not to attempt to modify the virtualenv. If you omit this flag, Pants will create a standalone, mutable virtualenv that includes `pip`, and that you can modify, but this method is slower.
 
-`./pants export` will also create a virtualenv for certain Python tools you use via Pants, like
-formatters like Black and Isort. This allows you to configure your editor to use the same version
-of the tool that Pants uses for workflows like formatting on save.
+### Tool virtualenvs
 
-To disable a certain tool, set its `export` option to `false`, e.g.:
+`./pants export` can also create a virtualenv for each of the Python tools you use via Pants, such as `black`, `isort`, `pytest`, `mypy`, `flake8` and so on (you can run `/pants help tools` to get a list of the tools Pants uses). Use the tool name as the resolve name argument to the `--resolve` flag. This allows you to configure your editor to use the same version of the tool as Pants does for workflows like formatting on save.
 
-```toml pants.toml
-[black]
-export = false
-```
 
 Generated code
 --------------

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

@@ -11,12 +11,13 @@
 from typing import Any, DefaultDict, Iterable, cast
 
 from pants.backend.python.subsystems.setup import PythonSetup
-from pants.backend.python.target_types import PythonResolveField
+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.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
 from pants.backend.python.util_rules.pex_from_targets import RequirementsPexRequest
+from pants.backend.python.util_rules.pex_requirements import EntireLockfile, Lockfile
 from pants.core.goals.export import (
     Export,
     ExportError,
@@ -56,6 +57,11 @@ def debug_hint(self) -> str | None:
         return self.resolve
 
 
+@dataclass(frozen=True)
+class _ExportVenvForResolveRequest(EngineAwareParameter):
+    resolve: str
+
+
 @union(in_scope_types=[EnvironmentName])
 class ExportPythonToolSentinel:
     """Python tools use this as an entry point to say how to export their tool virtualenv.
@@ -151,6 +157,7 @@ async def do_export(
             description,
             dest,
             post_processing_cmds=[PostProcessingCommand(["ln", "-s", venv_abspath, output_path])],
+            resolve=req.resolve_name or None,
         )
     else:
         # Note that an internal-only pex will always have the `python` field set.
@@ -196,8 +203,91 @@ async def do_export(
                 # Remove the requirements and pex pexes, to avoid confusion.
                 PostProcessingCommand(["rm", "-rf", tmpdir_under_digest_root]),
             ],
+            resolve=req.resolve_name or None,
+        )
+
+
+@dataclass(frozen=True)
+class MaybeExportResult:
+    result: ExportResult | None
+
+
+@rule
+async def export_virtualenv_for_resolve(
+    request: _ExportVenvForResolveRequest,
+    python_setup: PythonSetup,
+    union_membership: UnionMembership,
+) -> MaybeExportResult:
+    resolve = request.resolve
+    lockfile_path = python_setup.resolves.get(resolve)
+    if lockfile_path:
+        # It's a user resolve.
+        lockfile = Lockfile(
+            file_path=lockfile_path,
+            file_path_description_of_origin=f"the resolve `{resolve}`",
+            resolve_name=resolve,
+        )
+
+        interpreter_constraints = InterpreterConstraints(
+            python_setup.resolves_to_interpreter_constraints.get(
+                request.resolve, python_setup.interpreter_constraints
+            )
         )
 
+        pex_request = PexRequest(
+            description="chosen_resolve.name",
+            output_filename=f"{path_safe(resolve)}.pex",
+            internal_only=True,
+            requirements=EntireLockfile(lockfile),
+            interpreter_constraints=interpreter_constraints,
+            # Packed layout should lead to the best performance in this use case.
+            layout=PexLayout.PACKED,
+        )
+    else:
+        # It's a tool resolve.
+        # TODO: Can we simplify tool lockfiles to be more uniform with user lockfiles?
+        #  It's unclear if we will need the ExportPythonToolSentinel runaround once we
+        #  remove the older export codepath below. It would be nice to be able to go from
+        #  resolve name -> EntireLockfile, regardless of whether the resolve happened to be
+        #  a user lockfile or a tool lockfile. Currently we have to get all the ExportPythonTools
+        #  and then check for the resolve name.  But this is OK for now, as it lets us
+        #  move towards deprecating that other codepath.
+        tool_export_types = cast(
+            "Iterable[type[ExportPythonToolSentinel]]",
+            union_membership.get(ExportPythonToolSentinel),
+        )
+        all_export_tool_requests = await MultiGet(
+            Get(ExportPythonTool, ExportPythonToolSentinel, tool_export_type())
+            for tool_export_type in tool_export_types
+        )
+        export_tool_request = next(
+            (etr for etr in all_export_tool_requests if etr.resolve_name == resolve), None
+        )
+        if not export_tool_request:
+            # No such Python resolve or tool, but it may be a resolve for a different language/backend,
+            # so we let the core export goal sort out whether it's an error or not.
+            return MaybeExportResult(None)
+        if not export_tool_request.pex_request:
+            raise ExportError(
+                f"Requested an export of `{resolve}` but that tool's exports were disabled with "
+                f"the `export=false` option. The per-tool `export=false` options will soon be "
+                f"deprecated anyway, so we recommend removing `export=false` from your config file "
+                f"and switching to using `--resolve`."
+            )
+        pex_request = export_tool_request.pex_request
+
+    dest_prefix = os.path.join("python", "virtualenvs")
+    export_result = await Get(
+        ExportResult,
+        VenvExportRequest(
+            pex_request,
+            dest_prefix,
+            resolve,
+            qualify_path_with_python_version=True,
+        ),
+    )
+    return MaybeExportResult(export_result)
+
 
 @rule
 async def export_virtualenv_for_targets(
@@ -269,7 +359,18 @@ async def export_virtualenvs(
     python_setup: PythonSetup,
     dist_dir: DistDir,
     union_membership: UnionMembership,
+    export_subsys: ExportSubsystem,
 ) -> ExportResults:
+    if export_subsys.options.resolve:
+        if request.targets:
+            raise ExportError("If using the `--resolve` option, do not also provide target specs.")
+        maybe_venvs = await MultiGet(
+            Get(MaybeExportResult, _ExportVenvForResolveRequest(resolve))
+            for resolve in export_subsys.options.resolve
+        )
+        return ExportResults(mv.result for mv in maybe_venvs if mv.result is not None)
+
+    # TODO: Deprecate this entire codepath.
     resolve_to_root_targets: DefaultDict[str, list[Target]] = defaultdict(list)
     for tgt in request.targets:
         if not tgt.has_field(PythonResolveField):

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

@@ -10,6 +10,7 @@
 from pants.backend.python import target_types_rules
 from pants.backend.python.goals import export
 from pants.backend.python.goals.export import ExportVenvsRequest
+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.base.specs import RawSpecs, RecursiveGlobSpec
@@ -29,6 +30,7 @@ def rule_runner() -> RuleRunner:
             *pex_from_targets.rules(),
             *target_types_rules.rules(),
             *distdir.rules(),
+            *flake8_subsystem.rules(),
             QueryRule(Targets, [RawSpecs]),
             QueryRule(ExportResults, [ExportVenvsRequest]),
         ],
@@ -38,7 +40,7 @@ def rule_runner() -> RuleRunner:
 
 @pytest.mark.parametrize("enable_resolves", [False, True])
 @pytest.mark.parametrize("symlink", [False, True])
-def test_export_venv_pipified(
+def test_export_venv_old_codepath(
     rule_runner: RuleRunner,
     enable_resolves: bool,
     symlink: bool,
@@ -120,6 +122,90 @@ def test_export_venv_pipified(
         assert reldirs == [
             "python/virtualenvs/a",
             "python/virtualenvs/b",
+            "python/virtualenvs/tools/flake8",
         ]
     else:
-        assert reldirs == ["python/virtualenv"]
+        assert reldirs == ["python/virtualenv", "python/virtualenvs/tools/flake8"]
+
+
+@pytest.mark.parametrize("symlink", [False, True])
+def test_export_venv_new_codepath(
+    rule_runner: RuleRunner,
+    symlink: bool,
+) -> None:
+    # We know that the current interpreter exists on the system.
+    vinfo = sys.version_info
+    current_interpreter = f"{vinfo.major}.{vinfo.minor}.{vinfo.micro}"
+    rule_runner.write_files(
+        {
+            "src/foo/BUILD": dedent(
+                """\
+                python_requirement(name='req1', requirements=['ansicolors==1.1.8'], resolve='a')
+                python_requirement(name='req2', requirements=['ansicolors==1.1.8'], resolve='b')
+                """
+            ),
+            "lock.txt": "ansicolors==1.1.8",
+        }
+    )
+
+    symlink_flag = f"--{'' if symlink else 'no-'}export-symlink-python-virtualenv"
+    rule_runner.set_options(
+        [
+            f"--python-interpreter-constraints=['=={current_interpreter}']",
+            "--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",
+            symlink_flag,
+        ],
+        env_inherit={"PATH", "PYENV_ROOT"},
+    )
+    all_results = rule_runner.request(ExportResults, [ExportVenvsRequest(targets=())])
+
+    for result, resolve in zip(all_results, ["a", "b", "flake8"]):
+        if symlink:
+            assert len(result.post_processing_cmds) == 1
+            ppc0 = result.post_processing_cmds[0]
+            assert ppc0.argv[0:2] == ("ln", "-s")
+            # The third arg is the full path to the venv under the pex_root, which we
+            # don't easily know here, so we ignore it in this comparison.
+            assert ppc0.argv[3] == os.path.join("{digest_root}", current_interpreter)
+            assert ppc0.extra_env == FrozenDict()
+        else:
+            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
+            # don't easily know here, so we ignore it in this comparison.
+
+            # The second arg is expected to be tmpdir/./pex.
+            tmpdir, pex_pex_name = os.path.split(os.path.normpath(ppc0.argv[1]))
+            assert pex_pex_name == "pex"
+            assert re.match(r"\{digest_root\}/\.[0-9a-f]{32}\.tmp", tmpdir)
+
+            # The third arg is expected to be tmpdir/{resolve}.pex.
+            req_pex_dir, req_pex_name = os.path.split(ppc0.argv[2])
+            assert req_pex_dir == tmpdir
+            assert req_pex_name == f"{resolve}.pex"
+
+            assert ppc0.argv[3:] == (
+                "venv",
+                "--pip",
+                "--collisions-ok",
+                f"{{digest_root}}/{current_interpreter}",
+            )
+            assert ppc0.extra_env["PEX_MODULE"] == "pex.tools"
+            assert ppc0.extra_env.get("PEX_ROOT") is not None
+
+            ppc1 = result.post_processing_cmds[1]
+            assert ppc1.argv == ("rm", "-rf", tmpdir)
+            assert ppc1.extra_env == FrozenDict()
+
+    reldirs = [result.reldir for result in all_results]
+    assert reldirs == [
+        "python/virtualenvs/a",
+        "python/virtualenvs/b",
+        "python/virtualenvs/flake8",
+    ]