Repository: pantsbuild/pants

Issue:
Running `export` for specific Python tools (or lockfiles) only
**Is your feature request related to a problem? Please describe.**
Sometimes I modify requirements for a specific Python tool, but I need to re-export all of the virtualenvs including other Python tools and the default one (`python-default`).
It takes quite some time.

**Describe the solution you'd like**
I'd like to re-run `./pants export` for a specific lockfile only.
(e.g., `./pants export --resolve=flake8,mypy`)


Discussion:

Resolving patch:
diff --git a/src/python/pants/backend/python/goals/export.py b/src/python/pants/backend/python/goals/export.py
--- a/src/python/pants/backend/python/goals/export.py
+++ b/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):
diff --git a/src/python/pants/core/goals/export.py b/src/python/pants/core/goals/export.py
--- a/src/python/pants/core/goals/export.py
+++ b/src/python/pants/core/goals/export.py
@@ -3,11 +3,18 @@
 
 from __future__ import annotations
 
+import itertools
 import os
 from dataclasses import dataclass
 from typing import Iterable, Mapping, Sequence, cast
 
 from pants.base.build_root import BuildRoot
+from pants.core.goals.generate_lockfiles import (
+    GenerateToolLockfileSentinel,
+    KnownUserResolveNames,
+    KnownUserResolveNamesRequest,
+    UnrecognizedResolveNamesError,
+)
 from pants.core.util_rules.distdir import DistDir
 from pants.core.util_rules.environments import _warn_on_non_local_environments
 from pants.engine.collection import Collection
@@ -21,6 +28,7 @@
 from pants.engine.rules import collect_rules, goal_rule
 from pants.engine.target import FilteredTargets, Target
 from pants.engine.unions import UnionMembership, union
+from pants.option.option_types import StrListOption
 from pants.util.dirutil import safe_rmtree
 from pants.util.frozendict import FrozenDict
 from pants.util.meta import frozen_after_init
@@ -73,6 +81,9 @@ class ExportResult:
     digest: Digest
     # Run these commands as local processes after the digest is materialized.
     post_processing_cmds: tuple[PostProcessingCommand, ...]
+    # Set for the common special case of exporting a resolve, and names that resolve.
+    # Set to None for other export results.
+    resolve: str | None
 
     def __init__(
         self,
@@ -81,11 +92,13 @@ def __init__(
         *,
         digest: Digest = EMPTY_DIGEST,
         post_processing_cmds: Iterable[PostProcessingCommand] = tuple(),
+        resolve: str | None = None,
     ):
         self.description = description
         self.reldir = reldir
         self.digest = digest
         self.post_processing_cmds = tuple(post_processing_cmds)
+        self.resolve = resolve
 
 
 class ExportResults(Collection[ExportResult]):
@@ -96,6 +109,18 @@ class ExportSubsystem(GoalSubsystem):
     name = "export"
     help = "Export Pants data for use in other tools, such as IDEs."
 
+    # NB: Only options that are relevant across many/most backends and languages
+    #  should be defined here.  Backend-specific options should be defined in that backend
+    #  as plugin options on this subsystem.
+
+    # Exporting resolves is a common use-case for `export`, often the primary one, so we
+    # add affordances for it at the core goal level.
+    resolve = StrListOption(
+        default=[],
+        help="Export the specified resolve(s). The export format is backend-specific, "
+        "e.g., Python resolves are exported as virtualenvs.",
+    )
+
 
 class Export(Goal):
     subsystem_cls = ExportSubsystem
@@ -110,6 +135,7 @@ async def export(
     union_membership: UnionMembership,
     build_root: BuildRoot,
     dist_dir: DistDir,
+    export_subsys: ExportSubsystem,
 ) -> Export:
     request_types = cast("Iterable[type[ExportRequest]]", union_membership.get(ExportRequest))
     requests = tuple(request_type(targets) for request_type in request_types)
@@ -127,6 +153,7 @@ async def export(
     dist_digest = await Get(Digest, AddPrefix(merged_digest, output_dir))
     workspace.write_digest(dist_digest)
     environment = await Get(EnvironmentVars, EnvironmentVarsRequest(["PATH"]))
+    resolves_exported = set()
     for result in flattened_results:
         result_dir = os.path.join(output_dir, result.reldir)
         digest_root = os.path.join(build_root.path, result_dir)
@@ -140,8 +167,31 @@ async def export(
             ipr = await Effect(InteractiveProcessResult, InteractiveProcess, ip)
             if ipr.exit_code:
                 raise ExportError(f"Failed to write {result.description} to {result_dir}")
-
+        if result.resolve:
+            resolves_exported.add(result.resolve)
         console.print_stdout(f"Wrote {result.description} to {result_dir}")
+
+    unexported_resolves = sorted((set(export_subsys.resolve) - resolves_exported))
+    if unexported_resolves:
+        all_known_user_resolve_names = await MultiGet(
+            Get(KnownUserResolveNames, KnownUserResolveNamesRequest, request())
+            for request in union_membership.get(KnownUserResolveNamesRequest)
+        )
+        all_valid_resolve_names = sorted(
+            [
+                *itertools.chain.from_iterable(kurn.names for kurn in all_known_user_resolve_names),
+                *(
+                    sentinel.resolve_name
+                    for sentinel in union_membership.get(GenerateToolLockfileSentinel)
+                ),
+            ]
+        )
+        raise UnrecognizedResolveNamesError(
+            unexported_resolves,
+            all_valid_resolve_names,
+            description_of_origin="the option --export-resolve",
+        )
+
     return Export(exit_code=0)
 
 
diff --git a/src/python/pants/core/goals/generate_lockfiles.py b/src/python/pants/core/goals/generate_lockfiles.py
--- a/src/python/pants/core/goals/generate_lockfiles.py
+++ b/src/python/pants/core/goals/generate_lockfiles.py
@@ -337,10 +337,9 @@ def activated(cls, union_membership: UnionMembership) -> bool:
             Only generate lockfiles for the specified resolve(s).
 
             Resolves are the logical names for the different lockfiles used in your project.
-            For your own code's dependencies, these come from the option
-            `[python].resolves`. For tool lockfiles, resolve
-            names are the options scope for that tool such as `black`, `pytest`, and
-            `mypy-protobuf`.
+            For your own code's dependencies, these come from backend-specific configuration
+            such as `[python].resolves`. For tool lockfiles, resolve names are the options
+            scope for that tool such as `black`, `pytest`, and `mypy-protobuf`.
 
             For example, you can run `{bin_name()} generate-lockfiles --resolve=black
             --resolve=pytest --resolve=data-science` to only generate lockfiles for those


Test patch:
diff --git a/src/python/pants/backend/python/goals/export_test.py b/src/python/pants/backend/python/goals/export_test.py
--- a/src/python/pants/backend/python/goals/export_test.py
+++ b/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",
+    ]
diff --git a/src/python/pants/core/goals/export_test.py b/src/python/pants/core/goals/export_test.py
--- a/src/python/pants/core/goals/export_test.py
+++ b/src/python/pants/core/goals/export_test.py
@@ -16,9 +16,11 @@
     ExportRequest,
     ExportResult,
     ExportResults,
+    ExportSubsystem,
     PostProcessingCommand,
     export,
 )
+from pants.core.goals.generate_lockfiles import KnownUserResolveNames, KnownUserResolveNamesRequest
 from pants.core.util_rules.distdir import DistDir
 from pants.core.util_rules.environments import (
     EnvironmentField,
@@ -34,7 +36,7 @@
 from pants.engine.rules import QueryRule
 from pants.engine.target import Target, Targets
 from pants.engine.unions import UnionMembership, UnionRule
-from pants.testutil.option_util import create_options_bootstrapper
+from pants.testutil.option_util import create_options_bootstrapper, create_subsystem
 from pants.testutil.rule_runner import (
     MockEffect,
     MockGet,
@@ -99,6 +101,7 @@ def run_export_rule(rule_runner: RuleRunner, targets: List[Target]) -> Tuple[int
                 union_membership,
                 BuildRoot(),
                 DistDir(relpath=Path("dist")),
+                create_subsystem(ExportSubsystem, resolve=[]),
             ],
             mock_gets=[
                 MockGet(
@@ -129,6 +132,7 @@ def run_export_rule(rule_runner: RuleRunner, targets: List[Target]) -> Tuple[int
                 rule_runner.do_not_use_mock(Digest, (MergeDigests,)),
                 rule_runner.do_not_use_mock(Digest, (AddPrefix,)),
                 rule_runner.do_not_use_mock(EnvironmentVars, (EnvironmentVarsRequest,)),
+                rule_runner.do_not_use_mock(KnownUserResolveNames, (KnownUserResolveNamesRequest,)),
                 MockEffect(
                     output_type=InteractiveProcessResult,
                     input_types=(InteractiveProcess,),
@@ -230,6 +234,7 @@ def test_warnings_for_non_local_target_environments(
                 union_membership,
                 BuildRoot(),
                 DistDir(relpath=Path("dist")),
+                create_subsystem(ExportSubsystem, resolve=[]),
             ],
             mock_gets=[
                 MockGet(
@@ -253,6 +258,7 @@ def test_warnings_for_non_local_target_environments(
                 ),
                 rule_runner.do_not_use_mock(Digest, (AddPrefix,)),
                 rule_runner.do_not_use_mock(EnvironmentVars, (EnvironmentVarsRequest,)),
+                rule_runner.do_not_use_mock(KnownUserResolveNames, (KnownUserResolveNamesRequest,)),
                 MockEffect(
                     output_type=InteractiveProcessResult,
                     input_types=(InteractiveProcess,),