[internal] Add experimental_compatible_resolves to python_requirement and use it to generate lockfiles

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

@@ -24,9 +24,8 @@
 from pants.backend.python.subsystems.setup import PythonSetup
 from pants.backend.python.target_types import (
     EntryPoint,
-    InterpreterConstraintsField,
+    PythonCompatibleResolvesField,
     PythonRequirementsField,
-    PythonResolveField,
     UnrecognizedResolveNamesError,
 )
 from pants.backend.python.util_rules.interpreter_constraints import InterpreterConstraints
@@ -47,7 +46,7 @@
 from pants.engine.goal import Goal, GoalSubsystem
 from pants.engine.process import ProcessCacheScope, ProcessResult
 from pants.engine.rules import Get, MultiGet, collect_rules, goal_rule, rule
-from pants.engine.target import AllTargets, TransitiveTargets, TransitiveTargetsRequest
+from pants.engine.target import AllTargets
 from pants.engine.unions import UnionMembership, union
 from pants.util.logging import LogLevel
 from pants.util.ordered_set import FrozenOrderedSet
@@ -272,51 +271,30 @@ async def setup_user_lockfile_requests(
     if not python_setup.enable_resolves:
         return _UserLockfileRequests()
 
-    # First, associate all resolves with their consumers.
-    resolves_to_roots = defaultdict(list)
+    resolve_to_requirements_fields = defaultdict(set)
     for tgt in all_targets:
-        if not tgt.has_field(PythonResolveField):
+        if not tgt.has_field(PythonCompatibleResolvesField):
             continue
-        tgt[PythonResolveField].validate(python_setup)
-        resolve = tgt[PythonResolveField].value
-        if resolve is None:
-            continue
-        resolves_to_roots[resolve].append(tgt.address)
+        tgt[PythonCompatibleResolvesField].validate(python_setup)
+        for resolve in tgt[PythonCompatibleResolvesField].value_or_default(python_setup):
+            resolve_to_requirements_fields[resolve].add(tgt[PythonRequirementsField])
 
-    # Expand the resolves for all specified.
-    transitive_targets_per_resolve = await MultiGet(
-        Get(TransitiveTargets, TransitiveTargetsRequest(resolves_to_roots[resolve]))
-        for resolve in requested
-    )
-    pex_requirements_per_resolve = []
-    interpreter_constraints_per_resolve = []
-    for transitive_targets in transitive_targets_per_resolve:
-        req_fields = []
-        ic_fields = []
-        for tgt in transitive_targets.closure:
-            if tgt.has_field(PythonRequirementsField):
-                req_fields.append(tgt[PythonRequirementsField])
-            if tgt.has_field(InterpreterConstraintsField):
-                ic_fields.append(tgt[InterpreterConstraintsField])
-        pex_requirements_per_resolve.append(
-            PexRequirements.create_from_requirement_fields(req_fields)
-        )
-        interpreter_constraints_per_resolve.append(
-            InterpreterConstraints.create_from_compatibility_fields(ic_fields, python_setup)
-        )
+    # TODO: Figure out how to determine which interpreter constraints to use for each resolve...
+    #  Note that `python_requirement` does not have interpreter constraints, so we either need to
+    #  inspect all consumers of that resolve or start to closely couple the resolve with the
+    #  interpreter constraints (a "context").
 
-    requests = (
+    return _UserLockfileRequests(
         PythonLockfileRequest(
-            requirements.req_strings,
-            interpreter_constraints,
+            PexRequirements.create_from_requirement_fields(
+                resolve_to_requirements_fields[resolve]
+            ).req_strings,
+            InterpreterConstraints(python_setup.interpreter_constraints),
             resolve_name=resolve,
             lockfile_dest=python_setup.resolves[resolve],
         )
-        for resolve, requirements, interpreter_constraints in zip(
-            requested, pex_requirements_per_resolve, interpreter_constraints_per_resolve
-        )
+        for resolve in requested
     )
-    return _UserLockfileRequests(requests)
 
 
 # --------------------------------------------------------------------------------------

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

@@ -3,18 +3,26 @@
 
 from __future__ import annotations
 
+from textwrap import dedent
+
 import pytest
 
 from pants.backend.python.goals.lockfile import (
     AmbiguousResolveNamesError,
     PythonLockfileRequest,
     PythonToolLockfileSentinel,
+    _SpecifiedUserResolves,
+    _UserLockfileRequests,
     determine_resolves_to_generate,
     filter_tool_lockfile_requests,
+    setup_user_lockfile_requests,
 )
 from pants.backend.python.subsystems.python_tool_base import DEFAULT_TOOL_LOCKFILE, NO_TOOL_LOCKFILE
-from pants.backend.python.target_types import UnrecognizedResolveNamesError
+from pants.backend.python.subsystems.setup import PythonSetup
+from pants.backend.python.target_types import PythonRequirementTarget, UnrecognizedResolveNamesError
 from pants.backend.python.util_rules.interpreter_constraints import InterpreterConstraints
+from pants.engine.rules import SubsystemRule
+from pants.testutil.rule_runner import PYTHON_BOOTSTRAP_ENV, QueryRule, RuleRunner
 from pants.util.ordered_set import FrozenOrderedSet
 
 
@@ -111,3 +119,59 @@ def assert_filtered(
     assert f"`[{default_tool.resolve_name}].lockfile` is set to `{DEFAULT_TOOL_LOCKFILE}`" in str(
         exc.value
     )
+
+
+def test_multiple_resolves() -> None:
+    rule_runner = RuleRunner(
+        rules=[
+            setup_user_lockfile_requests,
+            SubsystemRule(PythonSetup),
+            QueryRule(_UserLockfileRequests, [_SpecifiedUserResolves]),
+        ],
+        target_types=[PythonRequirementTarget],
+    )
+    rule_runner.write_files(
+        {
+            "BUILD": dedent(
+                """\
+                python_requirement(
+                    name='both',
+                    requirements=['both1', 'both2'],
+                    experimental_compatible_resolves=['a', 'b'],
+                )
+                python_requirement(
+                    name='a',
+                    requirements=['a'],
+                    experimental_compatible_resolves=['a'],
+                )
+                python_requirement(
+                    name='b',
+                    requirements=['b'],
+                    experimental_compatible_resolves=['b'],
+                )
+                """
+            ),
+        }
+    )
+    rule_runner.set_options(
+        [
+            "--python-experimental-resolves={'a': 'a.lock', 'b': 'b.lock'}",
+            "--python-enable-resolves",
+        ],
+        env_inherit=PYTHON_BOOTSTRAP_ENV,
+    )
+    result = rule_runner.request(_UserLockfileRequests, [_SpecifiedUserResolves(["a", "b"])])
+    assert set(result) == {
+        PythonLockfileRequest(
+            FrozenOrderedSet(["a", "both1", "both2"]),
+            InterpreterConstraints(PythonSetup.default_interpreter_constraints),
+            resolve_name="a",
+            lockfile_dest="a.lock",
+        ),
+        PythonLockfileRequest(
+            FrozenOrderedSet(["b", "both1", "both2"]),
+            InterpreterConstraints(PythonSetup.default_interpreter_constraints),
+            resolve_name="b",
+            lockfile_dest="b.lock",
+        ),
+    }

src/python/pants/backend/python/target_types.py

@@ -125,6 +125,7 @@ def __init__(
 
 class PythonResolveField(StringField, AsyncFieldMixin):
     alias = "experimental_resolve"
+    required = False
     help = (
         "The resolve from `[python].experimental_resolves` to use.\n\n"
         "If not defined, will default to `[python].default_resolve`.\n\n"
@@ -158,6 +159,32 @@ def resolve_and_lockfile(self, python_setup: PythonSetup) -> tuple[str, str] | N
         return (resolve, python_setup.resolves[resolve])
 
 
+class PythonCompatibleResolvesField(StringSequenceField, AsyncFieldMixin):
+    alias = "experimental_compatible_resolves"
+    required = False
+    help = (
+        "The set of resolves from `[python].experimental_resolves` that this target is "
+        "compatible with.\n\n"
+        "If not defined, will default to `[python].default_resolve`.\n\n"
+        "Only applies if `[python].enable_resolves` is true.\n\n"
+        "This field is experimental and may change without the normal deprecation policy."
+        # TODO: Document expectations for dependencies once we validate that.
+    )
+
+    def value_or_default(self, python_setup: PythonSetup) -> tuple[str, ...]:
+        return self.value or (python_setup.default_resolve,)
+
+    def validate(self, python_setup: PythonSetup) -> None:
+        """Check that the resolve names are recognized."""
+        invalid_resolves = set(self.value_or_default(python_setup)) - set(python_setup.resolves)
+        if invalid_resolves:
+            raise UnrecognizedResolveNamesError(
+                sorted(invalid_resolves),
+                python_setup.resolves.keys(),
+                description_of_origin=f"the field `{self.alias}` in the target {self.address}",
+            )
+
+
 # -----------------------------------------------------------------------------------------------
 # `pex_binary` target
 # -----------------------------------------------------------------------------------------------
@@ -908,6 +935,21 @@ def normalize_module_mapping(
     return FrozenDict({canonicalize_project_name(k): tuple(v) for k, v in (mapping or {}).items()})
 
 
+class PythonRequirementCompatibleResolvesField(PythonCompatibleResolvesField):
+    help = (
+        "The resolves from `[python].experimental_resolves` that this requirement should be "
+        "included in.\n\n"
+        "If not defined, will default to `[python].default_resolve`.\n\n"
+        "When generating a lockfile for a particular resolve via the `generate-lockfiles` goal, "
+        "it will include all requirements that are declared compatible with that resolve. "
+        "First-party targets like `python_source` and `pex_binary` then declare which resolve(s) "
+        "they use via the `experimental_resolve` and `experimental_compatible_resolves` field; so, "
+        "for your first-party code to use a particular `python_requirement` target, that "
+        "requirement must be included in the resolve(s) "
+        "used by that code."
+    )
+
+
 class PythonRequirementTarget(Target):
     alias = "python_requirement"
     core_fields = (
@@ -916,6 +958,7 @@ class PythonRequirementTarget(Target):
         PythonRequirementsField,
         PythonRequirementModulesField,
         PythonRequirementTypeStubModulesField,
+        PythonCompatibleResolvesField,
     )
     help = (
         "A Python requirement installable by pip.\n\n"

src/python/pants/jvm/target_types.py

@@ -154,9 +154,8 @@ class JvmProvidesTypesField(StringSequenceField):
 
 class JvmArtifactCompatibleResolvesField(JvmCompatibleResolvesField):
     help = (
-        "The resolves that this artifact should be included in.\n\n"
+        "The resolves from `[jvm].resolves` that this artifact should be included in.\n\n"
         "If not defined, will default to `[jvm].default_resolve`.\n\n"
-        "Each name must be defined as a resolve in `[jvm].resolves`.\n\n"
         "When generating a lockfile for a particular resolve via the `coursier-resolve` goal, "
         "it will include all artifacts that are declared compatible with that resolve. First-party "
         "targets like `java_source` and `scala_source` then declare which resolve(s) they use "