Improve export goal to handle multiple Python resolves

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

@@ -3,50 +3,83 @@
 
 from __future__ import annotations
 
+import logging
 import os
+from collections import defaultdict
 from dataclasses import dataclass
+from typing import DefaultDict
 
 from pants.backend.python.subsystems.setup import PythonSetup
+from pants.backend.python.target_types import PythonResolveField
 from pants.backend.python.util_rules.interpreter_constraints import InterpreterConstraints
 from pants.backend.python.util_rules.pex import VenvPex, VenvPexProcess
 from pants.backend.python.util_rules.pex_environment import PexEnvironment
 from pants.backend.python.util_rules.pex_from_targets import RequirementsPexRequest
-from pants.core.goals.export import ExportableData, ExportableDataRequest, ExportError, Symlink
-from pants.engine.internals.selectors import Get
+from pants.core.goals.export import ExportError, ExportRequest, ExportResult, ExportResults, Symlink
+from pants.core.util_rules.distdir import DistDir
+from pants.engine.engine_aware import EngineAwareParameter
+from pants.engine.internals.selectors import Get, MultiGet
 from pants.engine.process import ProcessResult
 from pants.engine.rules import collect_rules, rule
+from pants.engine.target import Target
 from pants.engine.unions import UnionRule
+from pants.util.docutil import bin_name
+from pants.util.strutil import path_safe
+
+logger = logging.getLogger(__name__)
 
 
 @dataclass(frozen=True)
-class ExportedVenvRequest(ExportableDataRequest):
+class ExportVenvsRequest(ExportRequest):
     pass
 
 
+@dataclass(frozen=True)
+class _ExportVenvRequest(EngineAwareParameter):
+    resolve: str | None
+    root_python_targets: tuple[Target, ...]
+
+    def debug_hint(self) -> str | None:
+        return self.resolve
+
+
 @rule
-async def export_venv(
-    request: ExportedVenvRequest, python_setup: PythonSetup, pex_env: PexEnvironment
-) -> ExportableData:
-    # Pick a single interpreter for the venv.
-    interpreter_constraints = InterpreterConstraints.create_from_targets(
-        request.targets, python_setup
-    )
-    if not interpreter_constraints:
-        # If there were no targets that defined any constraints, fall back to the global ones.
-        interpreter_constraints = InterpreterConstraints(python_setup.interpreter_constraints)
+async def export_virtualenv(
+    request: _ExportVenvRequest, python_setup: PythonSetup, pex_env: PexEnvironment
+) -> ExportResult:
+    if request.resolve:
+        interpreter_constraints = InterpreterConstraints(
+            python_setup.resolves_to_interpreter_constraints.get(
+                request.resolve, python_setup.interpreter_constraints
+            )
+        )
+    else:
+        interpreter_constraints = InterpreterConstraints.create_from_targets(
+            request.root_python_targets, python_setup
+        ) or InterpreterConstraints(python_setup.interpreter_constraints)
+
     min_interpreter = interpreter_constraints.snap_to_minimum(python_setup.interpreter_universe)
     if not min_interpreter:
-        raise ExportError(
-            "The following interpreter constraints were computed for all the targets for which "
-            f"export was requested: {interpreter_constraints}. There is no python interpreter "
-            "compatible with these constraints. Please restrict the target set to one that shares "
-            "a compatible interpreter."
+        err_msg = (
+            (
+                f"The resolve '{request.resolve}' (from `[python].resolves`) has invalid interpreter "
+                f"constraints, which are set via `[python].resolves_to_interpreter_constraints`: "
+                f"{interpreter_constraints}. Could not determine the minimum compatible interpreter."
+            )
+            if request.resolve
+            else (
+                "The following interpreter constraints were computed for all the targets for which "
+                f"export was requested: {interpreter_constraints}. There is no python interpreter "
+                "compatible with these constraints. Please restrict the target set to one that shares "
+                "a compatible interpreter."
+            )
         )
+        raise ExportError(err_msg)
 
     venv_pex = await Get(
         VenvPex,
         RequirementsPexRequest(
-            (tgt.address for tgt in request.targets),
+            (tgt.address for tgt in request.root_python_targets),
             internal_only=True,
             hardcoded_interpreter_constraints=min_interpreter,
         ),
@@ -68,15 +101,52 @@ async def export_venv(
     )
     py_version = res.stdout.strip().decode()
 
-    return ExportableData(
-        f"virtualenv for {min_interpreter}",
-        os.path.join("python", "virtualenv"),
+    dest = (
+        os.path.join("python", "virtualenvs", path_safe(request.resolve))
+        if request.resolve
+        else os.path.join("python", "virtualenv")
+    )
+    return ExportResult(
+        f"virtualenv for the resolve '{request.resolve}' (using {min_interpreter})",
+        dest,
         symlinks=[Symlink(venv_abspath, py_version)],
     )
 
 
+@rule
+async def export_virtualenvs(
+    request: ExportVenvsRequest, python_setup: PythonSetup, dist_dir: DistDir
+) -> ExportResults:
+    resolve_to_root_targets: DefaultDict[str, list[Target]] = defaultdict(list)
+    for tgt in request.targets:
+        if not tgt.has_field(PythonResolveField):
+            continue
+        resolve = tgt[PythonResolveField].normalized_value(python_setup)
+        resolve_to_root_targets[resolve].append(tgt)
+
+    venvs = await MultiGet(
+        Get(
+            ExportResult,
+            _ExportVenvRequest(resolve if python_setup.enable_resolves else None, tuple(tgts)),
+        )
+        for resolve, tgts in resolve_to_root_targets.items()
+    )
+
+    no_resolves_dest = dist_dir.relpath / "python" / "virtualenv"
+    if venvs and python_setup.enable_resolves and no_resolves_dest.exists():
+        logger.warning(
+            f"Because `[python].enable_resolves` is true, `{bin_name()} export ::` no longer "
+            f"writes virtualenvs to {no_resolves_dest}, but instead underneath "
+            f"{dist_dir.relpath / 'python' / 'virtualenvs'}. You will need to "
+            "update your IDE to point to the new virtualenv.\n\n"
+            f"To silence this error, delete {no_resolves_dest}"
+        )
+
+    return ExportResults(venvs)
+
+
 def rules():
     return [
         *collect_rules(),
-        UnionRule(ExportableDataRequest, ExportedVenvRequest),
+        UnionRule(ExportRequest, ExportVenvsRequest),
     ]

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

@@ -1,16 +1,18 @@
 # Copyright 2021 Pants project contributors (see CONTRIBUTORS.md).
 # Licensed under the Apache License, Version 2.0 (see LICENSE).
 import sys
+from textwrap import dedent
 
 import pytest
 
 from pants.backend.python import target_types_rules
 from pants.backend.python.goals import export
-from pants.backend.python.goals.export import ExportedVenvRequest
+from pants.backend.python.goals.export import ExportVenvsRequest
 from pants.backend.python.target_types import PythonRequirementTarget
 from pants.backend.python.util_rules import pex_from_targets
 from pants.base.specs import AddressSpecs, DescendantAddresses
-from pants.core.goals.export import ExportableData
+from pants.core.goals.export import ExportResults
+from pants.core.util_rules import distdir
 from pants.engine.rules import QueryRule
 from pants.engine.target import Targets
 from pants.testutil.rule_runner import RuleRunner
@@ -23,28 +25,59 @@ def rule_runner() -> RuleRunner:
             *export.rules(),
             *pex_from_targets.rules(),
             *target_types_rules.rules(),
+            *distdir.rules(),
             QueryRule(Targets, [AddressSpecs]),
-            QueryRule(ExportableData, [ExportedVenvRequest]),
+            QueryRule(ExportResults, [ExportVenvsRequest]),
         ],
         target_types=[PythonRequirementTarget],
     )
 
 
-def test_export_venv(rule_runner: RuleRunner) -> None:
+def test_export_venvs(rule_runner: RuleRunner) -> 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.set_options(
-        [f"--python-interpreter-constraints=['=={current_interpreter}']"],
-        env_inherit={"PATH", "PYENV_ROOT"},
-    )
     rule_runner.write_files(
-        {"src/foo/BUILD": "python_requirement(name='req', requirements=['ansicolors==1.1.8'])"}
+        {
+            "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",
+        }
     )
-    targets = rule_runner.request(Targets, [AddressSpecs([DescendantAddresses("src/foo")])])
-    data = rule_runner.request(ExportableData, [ExportedVenvRequest(targets)])
-    assert len(data.symlinks) == 1
-    symlink = data.symlinks[0]
-    assert symlink.link_rel_path == current_interpreter
-    assert "named_caches/pex_root/venvs/" in symlink.source_path
+
+    def run(enable_resolves: bool) -> ExportResults:
+        rule_runner.set_options(
+            [
+                f"--python-interpreter-constraints=['=={current_interpreter}']",
+                "--python-resolves={'a': 'lock.txt', 'b': 'lock.txt'}",
+                f"--python-enable-resolves={enable_resolves}",
+                # Turn off lockfile validation to make the test simpler.
+                "--python-invalid-lockfile-behavior=ignore",
+            ],
+            env_inherit={"PATH", "PYENV_ROOT"},
+        )
+        targets = rule_runner.request(Targets, [AddressSpecs([DescendantAddresses("src/foo")])])
+        all_results = rule_runner.request(ExportResults, [ExportVenvsRequest(targets)])
+
+        for result in all_results:
+            assert len(result.symlinks) == 1
+            symlink = result.symlinks[0]
+            assert symlink.link_rel_path == current_interpreter
+            assert "named_caches/pex_root/venvs/" in symlink.source_path
+
+        return all_results
+
+    resolve_results = run(enable_resolves=True)
+    assert len(resolve_results) == 2
+    assert {result.reldir for result in resolve_results} == {
+        "python/virtualenvs/a",
+        "python/virtualenvs/b",
+    }
+
+    no_resolve_results = run(enable_resolves=False)
+    assert len(no_resolve_results) == 1
+    assert no_resolve_results[0].reldir == "python/virtualenv"

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

@@ -9,6 +9,7 @@
 
 from pants.base.build_root import BuildRoot
 from pants.core.util_rules.distdir import DistDir
+from pants.engine.collection import Collection
 from pants.engine.console import Console
 from pants.engine.fs import EMPTY_DIGEST, AddPrefix, Digest, MergeDigests, Workspace
 from pants.engine.goal import Goal, GoalSubsystem
@@ -26,7 +27,7 @@ class ExportError(Exception):
 
 @union
 @dataclass(frozen=True)
-class ExportableDataRequest:
+class ExportRequest:
     """A union for exportable data provided by a backend.
 
     Subclass and install a member of this type to export data.
@@ -41,7 +42,7 @@ class Symlink:
 
     source_path may be absolute, or relative to the repo root.
 
-    link_rel_path is relative to the enclosing ExportableData's reldir, and will be
+    link_rel_path is relative to the enclosing ExportResult's reldir, and will be
     absolutized when a location for that dir is chosen.
     """
 
@@ -51,7 +52,7 @@ class Symlink:
 
 @frozen_after_init
 @dataclass(unsafe_hash=True)
-class ExportableData:
+class ExportResult:
     description: str
     # Materialize digests and create symlinks under this reldir.
     reldir: str
@@ -75,6 +76,10 @@ def __init__(
         self.symlinks = tuple(symlinks)
 
 
+class ExportResults(Collection[ExportResult]):
+    pass
+
+
 class ExportSubsystem(GoalSubsystem):
     name = "export"
     help = "Export Pants data for use in other tools, such as IDEs."
@@ -88,35 +93,34 @@ class Export(Goal):
 async def export(
     console: Console,
     targets: Targets,
-    export_subsystem: ExportSubsystem,
     workspace: Workspace,
     union_membership: UnionMembership,
     build_root: BuildRoot,
     dist_dir: DistDir,
 ) -> Export:
-    request_types = cast(
-        "Iterable[type[ExportableDataRequest]]", union_membership.get(ExportableDataRequest)
-    )
+    request_types = cast("Iterable[type[ExportRequest]]", union_membership.get(ExportRequest))
     requests = tuple(request_type(targets) for request_type in request_types)
-    exportables = await MultiGet(
-        Get(ExportableData, ExportableDataRequest, request) for request in requests
-    )
+    all_results = await MultiGet(Get(ExportResults, ExportRequest, request) for request in requests)
+    flattened_results = [res for results in all_results for res in results]
+
     prefixed_digests = await MultiGet(
-        Get(Digest, AddPrefix(exp.digest, exp.reldir)) for exp in exportables
+        Get(Digest, AddPrefix(result.digest, result.reldir)) for result in flattened_results
     )
     output_dir = os.path.join(str(dist_dir.relpath), "export")
     merged_digest = await Get(Digest, MergeDigests(prefixed_digests))
     dist_digest = await Get(Digest, AddPrefix(merged_digest, output_dir))
     workspace.write_digest(dist_digest)
-    for exp in exportables:
-        for symlink in exp.symlinks:
+    for result in flattened_results:
+        for symlink in result.symlinks:
             # Note that if symlink.source_path is an abspath, join returns it unchanged.
             source_abspath = os.path.join(build_root.path, symlink.source_path)
             link_abspath = os.path.abspath(
-                os.path.join(output_dir, exp.reldir, symlink.link_rel_path)
+                os.path.join(output_dir, result.reldir, symlink.link_rel_path)
             )
             absolute_symlink(source_abspath, link_abspath)
-        console.print_stdout(f"Wrote {exp.description} to {os.path.join(output_dir, exp.reldir)}")
+        console.print_stdout(
+            f"Wrote {result.description} to {os.path.join(output_dir, result.reldir)}"
+        )
     return Export(exit_code=0)