A --py-resolve flag, to specify the resolves to export.

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,
@@ -35,7 +36,7 @@
 from pants.engine.rules import collect_rules, rule, rule_helper
 from pants.engine.target import Target
 from pants.engine.unions import UnionMembership, UnionRule, union
-from pants.option.option_types import BoolOption
+from pants.option.option_types import BoolOption, StrListOption
 from pants.util.docutil import bin_name
 from pants.util.strutil import path_safe, softwrap
 
@@ -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.
@@ -83,6 +89,8 @@ def debug_hint(self) -> str | None:
 
 
 class ExportPluginOptions:
+    py_resolve = StrListOption(default=[], help="Export virtualenvs for this resolve.")
+
     symlink_python_virtualenv = BoolOption(
         default=False,
         help="Export a symlink into a cached Python virtualenv.  This virtualenv will have no pip binary, "
@@ -199,6 +207,79 @@ async def do_export(
         )
 
 
+@rule
+async def export_virtualenv_for_resolve(
+    request: _ExportVenvForResolveRequest,
+    python_setup: PythonSetup,
+    union_membership: UnionMembership,
+) -> ExportResult:
+    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:
+            raise ExportError(f"No such resolve: {resolve}")
+        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."
+            )
+        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 export_result
+
+
 @rule
 async def export_virtualenv_for_targets(
     request: _ExportVenvRequest,
@@ -269,7 +350,20 @@ async def export_virtualenvs(
     python_setup: PythonSetup,
     dist_dir: DistDir,
     union_membership: UnionMembership,
+    export_subsys: ExportSubsystem,
 ) -> ExportResults:
+    if export_subsys.options.py_resolve:
+        if request.targets:
+            raise ExportError(
+                "If using export's --py-resolve option, do not also provide target specs."
+            )
+        venvs = await MultiGet(
+            Get(ExportResult, _ExportVenvForResolveRequest(resolve))
+            for resolve in export_subsys.options.py_resolve
+        )
+        return ExportResults(venvs)
+
+    # 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-py-resolve=a",
+            "--export-py-resolve=b",
+            "--export-py-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",
+    ]