Generify user lockfiles

This is pretty involved. We need to first use a union to discover all "known" resolves per language. We then need another union to convert all _requested_ resolves for that language into concrete LockfileRequest subclasses

# Rust tests and lints will be skipped. Delete if not intended.
[ci skip-rust]

# Building wheels and fs_util will be skipped. Delete if not intended.
[ci skip-build-wheels]

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

@@ -30,14 +30,17 @@
 )
 from pants.backend.python.util_rules.pex import PexRequest, PexRequirements, VenvPex, VenvPexProcess
 from pants.core.goals.generate_lockfiles import (
+    KnownUserResolveNames,
+    KnownUserResolveNamesRequest,
     Lockfile,
     LockfileRequest,
+    RequestedUserResolveNames,
     ToolLockfileSentinel,
+    UserLockfileRequests,
     WrappedLockfileRequest,
     determine_resolves_to_generate,
     filter_tool_lockfile_requests,
 )
-from pants.engine.collection import Collection
 from pants.engine.fs import (
     CreateDigest,
     Digest,
@@ -60,7 +63,7 @@
 class GenerateLockfilesSubsystem(GoalSubsystem):
     name = "generate-lockfiles"
     help = "Generate lockfiles for Python third-party dependencies."
-    required_union_implementations = (ToolLockfileSentinel,)
+    required_union_implementations = (ToolLockfileSentinel, KnownUserResolveNames)
 
     @classmethod
     def register_options(cls, register) -> None:
@@ -266,20 +269,31 @@ async def generate_lockfile(
     return Lockfile(final_lockfile_digest, req.resolve_name, req.lockfile_dest)
 
 
-class _SpecifiedUserResolves(Collection[str]):
+class RequestedPythonUserResolveNames(RequestedUserResolveNames):
     pass
 
 
-class _UserLockfileRequests(Collection[PythonLockfileRequest]):
+class KnownPythonUserResolveNamesRequest(KnownUserResolveNamesRequest):
     pass
 
 
+@rule
+def determine_python_user_resolves(
+    _: KnownPythonUserResolveNamesRequest, python_setup: PythonSetup
+) -> KnownUserResolveNames:
+    return KnownUserResolveNames(
+        names=tuple(python_setup.resolves.keys()),
+        option_name="[python].experimental_resolves",
+        requested_resolve_names_cls=RequestedPythonUserResolveNames,
+    )
+
+
 @rule
 async def setup_user_lockfile_requests(
-    requested: _SpecifiedUserResolves, all_targets: AllTargets, python_setup: PythonSetup
-) -> _UserLockfileRequests:
+    requested: RequestedPythonUserResolveNames, all_targets: AllTargets, python_setup: PythonSetup
+) -> UserLockfileRequests:
     if not python_setup.enable_resolves:
-        return _UserLockfileRequests()
+        return UserLockfileRequests()
 
     resolve_to_requirements_fields = defaultdict(set)
     for tgt in all_targets:
@@ -294,7 +308,7 @@ async def setup_user_lockfile_requests(
     #  inspect all consumers of that resolve or start to closely couple the resolve with the
     #  interpreter constraints (a "context").
 
-    return _UserLockfileRequests(
+    return UserLockfileRequests(
         PythonLockfileRequest(
             requirements=PexRequirements.create_from_requirement_fields(
                 resolve_to_requirements_fields[resolve],
@@ -317,20 +331,24 @@ async def generate_lockfiles_goal(
     workspace: Workspace,
     union_membership: UnionMembership,
     generate_lockfiles_subsystem: GenerateLockfilesSubsystem,
-    python_setup: PythonSetup,
 ) -> GenerateLockfilesGoal:
-    specified_user_resolves, specified_tool_sentinels = determine_resolves_to_generate(
-        python_setup.resolves.keys(),
-        union_membership[ToolLockfileSentinel],
-        generate_lockfiles_subsystem.resolve_names,
+    known_user_resolve_names = await MultiGet(
+        Get(KnownUserResolveNames, KnownUserResolveNamesRequest, request())
+        for request in union_membership.get(KnownUserResolveNamesRequest)
+    )
+    requested_user_resolve_names, requested_tool_sentinels = determine_resolves_to_generate(
+        known_user_resolve_names,
+        union_membership.get(ToolLockfileSentinel),
+        set(generate_lockfiles_subsystem.resolve_names),
     )
 
-    specified_user_requests = await Get(
-        _UserLockfileRequests, _SpecifiedUserResolves(specified_user_resolves)
+    all_specified_user_requests = await MultiGet(
+        Get(UserLockfileRequests, RequestedUserResolveNames, resolve_names)
+        for resolve_names in requested_user_resolve_names
     )
     specified_tool_requests = await MultiGet(
         Get(WrappedLockfileRequest, ToolLockfileSentinel, sentinel())
-        for sentinel in specified_tool_sentinels
+        for sentinel in requested_tool_sentinels
     )
     applicable_tool_requests = filter_tool_lockfile_requests(
         specified_tool_requests,
@@ -339,7 +357,10 @@ async def generate_lockfiles_goal(
 
     results = await MultiGet(
         Get(Lockfile, LockfileRequest, req)
-        for req in (*specified_user_requests, *applicable_tool_requests)
+        for req in (
+            *(req for reqs in all_specified_user_requests for req in reqs),
+            *applicable_tool_requests,
+        )
     )
 
     merged_digest = await Get(Digest, MergeDigests(res.digest for res in results))
@@ -351,4 +372,9 @@ async def generate_lockfiles_goal(
 
 
 def rules():
-    return (*collect_rules(), UnionRule(LockfileRequest, PythonLockfileRequest))
+    return (
+        *collect_rules(),
+        UnionRule(LockfileRequest, PythonLockfileRequest),
+        UnionRule(KnownUserResolveNamesRequest, KnownPythonUserResolveNamesRequest),
+        UnionRule(RequestedUserResolveNames, RequestedPythonUserResolveNames),
+    )

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

@@ -7,13 +7,13 @@
 
 from pants.backend.python.goals.lockfile import (
     PythonLockfileRequest,
-    _SpecifiedUserResolves,
-    _UserLockfileRequests,
+    RequestedPythonUserResolveNames,
     setup_user_lockfile_requests,
 )
 from pants.backend.python.subsystems.setup import PythonSetup
 from pants.backend.python.target_types import PythonRequirementTarget
 from pants.backend.python.util_rules.interpreter_constraints import InterpreterConstraints
+from pants.core.goals.generate_lockfiles import UserLockfileRequests
 from pants.engine.rules import SubsystemRule
 from pants.testutil.rule_runner import PYTHON_BOOTSTRAP_ENV, QueryRule, RuleRunner
 from pants.util.ordered_set import FrozenOrderedSet
@@ -24,7 +24,7 @@ def test_multiple_resolves() -> None:
         rules=[
             setup_user_lockfile_requests,
             SubsystemRule(PythonSetup),
-            QueryRule(_UserLockfileRequests, [_SpecifiedUserResolves]),
+            QueryRule(UserLockfileRequests, [RequestedPythonUserResolveNames]),
         ],
         target_types=[PythonRequirementTarget],
     )
@@ -58,7 +58,9 @@ def test_multiple_resolves() -> None:
         ],
         env_inherit=PYTHON_BOOTSTRAP_ENV,
     )
-    result = rule_runner.request(_UserLockfileRequests, [_SpecifiedUserResolves(["a", "b"])])
+    result = rule_runner.request(
+        UserLockfileRequests, [RequestedPythonUserResolveNames(["a", "b"])]
+    )
     assert set(result) == {
         PythonLockfileRequest(
             requirements=FrozenOrderedSet(["a", "both1", "both2"]),

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

@@ -3,9 +3,11 @@
 
 from __future__ import annotations
 
+import itertools
 from dataclasses import dataclass
 from typing import ClassVar, Iterable, Sequence
 
+from pants.engine.collection import Collection
 from pants.engine.fs import Digest
 from pants.engine.unions import union
 
@@ -56,6 +58,48 @@ class ToolLockfileSentinel:
     options_scope: ClassVar[str]
 
 
+class UserLockfileRequests(Collection[LockfileRequest]):
+    """All user resolves for a particular language ecosystem to build.
+
+    Each language ecosystem should set up a subclass of `RequestedUserResolveNames` (see its
+    docstring), and implement a rule going from that subclass -> UserLockfileRequests. Each element
+    in the returned `UserLockfileRequests` should be a subclass of `LockfileRequest`, like
+    `PythonLockfileRequest`.
+    """
+
+
+@union
+class KnownUserResolveNamesRequest:
+    """A hook for a language ecosystem to declare which resolves it has defined.
+
+    Each language ecosystem should set up a subclass and register it with a UnionRule. Implement a
+    rule that goes from the subclass -> KnownUserResolveNames, usually by simply reading the
+    `resolves` option from the relevant subsystem.
+    """
+
+
+@dataclass(frozen=True)
+class KnownUserResolveNames:
+    """All defined user resolves for a particular language ecosystem.
+
+    See KnownUserResolveNamesRequest for how to use this type. `option_name` should be formatted
+    like `[options-scope].resolves`
+    """
+
+    names: tuple[str, ...]
+    option_name: str
+    requested_resolve_names_cls: type[RequestedUserResolveNames]
+
+
+@union
+class RequestedUserResolveNames(Collection[str]):
+    """The user resolves requested for a particular language ecosystem.
+
+    Each language ecosystem should set up a subclass and register it with a UnionRule. Implement a
+    rule that goes from the subclass -> UserLockfileRequests.
+    """
+
+
 DEFAULT_TOOL_LOCKFILE = "<default>"
 NO_TOOL_LOCKFILE = "<none>"
 
@@ -100,10 +144,10 @@ def __init__(self, ambiguous_names: list[str]) -> None:
 
 
 def determine_resolves_to_generate(
-    all_user_resolves: Iterable[str],
+    all_known_user_resolve_names: Iterable[KnownUserResolveNames],
     all_tool_sentinels: Iterable[type[ToolLockfileSentinel]],
-    requested_resolve_names: Sequence[str],
-) -> tuple[list[str], list[type[ToolLockfileSentinel]]]:
+    requested_resolve_names: set[str],
+) -> tuple[list[RequestedUserResolveNames], list[type[ToolLockfileSentinel]]]:
     """Apply the `--resolve` option to determine which resolves are specified.
 
     Return a tuple of `(user_resolves, tool_lockfile_sentinels)`.
@@ -112,37 +156,45 @@ def determine_resolves_to_generate(
         sentinel.options_scope: sentinel for sentinel in all_tool_sentinels
     }
 
-    ambiguous_resolve_names = [
-        resolve_name
-        for resolve_name in all_user_resolves
-        if resolve_name in resolve_names_to_sentinels
-    ]
-    if ambiguous_resolve_names:
-        raise AmbiguousResolveNamesError(ambiguous_resolve_names)
+    # TODO: check for ambiguity: between tools and user resolves, and across distinct
+    #  `KnownUserResolveNames`s. Update AmbiguousResolveNamesError to say where the resolve
+    #  name is defined, whereas right now we hardcode it to be the `[python]` option.
 
     if not requested_resolve_names:
-        return list(all_user_resolves), list(all_tool_sentinels)
+        return [
+            known_resolve_names.requested_resolve_names_cls(known_resolve_names.names)
+            for known_resolve_names in all_known_user_resolve_names
+        ], list(all_tool_sentinels)
+
+    requested_user_resolve_names = []
+    for known_resolve_names in all_known_user_resolve_names:
+        requested = requested_resolve_names.intersection(known_resolve_names.names)
+        if requested:
+            requested_resolve_names -= requested
+            requested_user_resolve_names.append(
+                known_resolve_names.requested_resolve_names_cls(requested)
+            )
 
-    specified_user_resolves = []
     specified_sentinels = []
-    unrecognized_resolve_names = []
-    for resolve_name in requested_resolve_names:
-        sentinel = resolve_names_to_sentinels.get(resolve_name)
-        if sentinel:
+    for resolve, sentinel in resolve_names_to_sentinels.items():
+        if resolve in requested_resolve_names:
+            requested_resolve_names.discard(resolve)
             specified_sentinels.append(sentinel)
-        elif resolve_name in all_user_resolves:
-            specified_user_resolves.append(resolve_name)
-        else:
-            unrecognized_resolve_names.append(resolve_name)
 
-    if unrecognized_resolve_names:
+    if requested_resolve_names:
         raise UnrecognizedResolveNamesError(
-            unrecognized_resolve_names,
-            {*all_user_resolves, *resolve_names_to_sentinels.keys()},
+            unrecognized_resolve_names=sorted(requested_resolve_names),
+            all_valid_names={
+                *itertools.chain.from_iterable(
+                    known_resolve_names.names
+                    for known_resolve_names in all_known_user_resolve_names
+                ),
+                *resolve_names_to_sentinels.keys(),
+            },
             description_of_origin="the option `--generate-lockfiles-resolve`",
         )
 
-    return specified_user_resolves, specified_sentinels
+    return requested_user_resolve_names, specified_sentinels
 
 
 def filter_tool_lockfile_requests(

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

@@ -8,8 +8,9 @@
 from pants.core.goals.generate_lockfiles import (
     DEFAULT_TOOL_LOCKFILE,
     NO_TOOL_LOCKFILE,
-    AmbiguousResolveNamesError,
+    KnownUserResolveNames,
     LockfileRequest,
+    RequestedUserResolveNames,
     ToolLockfileSentinel,
     UnrecognizedResolveNamesError,
     WrappedLockfileRequest,
@@ -28,44 +29,60 @@ class Tool2(ToolLockfileSentinel):
     class Tool3(ToolLockfileSentinel):
         options_scope = "tool3"
 
-    all_user_resolves = ["u1", "u2", "u3"]
+    class Lang1Requested(RequestedUserResolveNames):
+        pass
+
+    class Lang2Requested(RequestedUserResolveNames):
+        pass
+
+    lang1_resolves = KnownUserResolveNames(
+        ("u1", "u2"), option_name="[lang1].resolves", requested_resolve_names_cls=Lang1Requested
+    )
+    lang2_resolves = KnownUserResolveNames(
+        ("u3",), option_name="[lang2].resolves", requested_resolve_names_cls=Lang2Requested
+    )
 
     def assert_chosen(
-        requested: list[str],
-        expected_user_resolves: list[str],
+        requested: set[str],
+        expected_user_resolves: list[RequestedUserResolveNames],
         expected_tools: list[type[ToolLockfileSentinel]],
     ) -> None:
         user_resolves, tools = determine_resolves_to_generate(
-            all_user_resolves, [Tool1, Tool2, Tool3], requested
+            [lang1_resolves, lang2_resolves], [Tool1, Tool2, Tool3], requested
         )
         assert user_resolves == expected_user_resolves
         assert tools == expected_tools
 
     assert_chosen(
-        [Tool2.options_scope, "u2"], expected_user_resolves=["u2"], expected_tools=[Tool2]
+        {Tool2.options_scope, "u2"},
+        expected_user_resolves=[Lang1Requested(["u2"])],
+        expected_tools=[Tool2],
     )
     assert_chosen(
-        [Tool1.options_scope, Tool3.options_scope],
+        {Tool1.options_scope, Tool3.options_scope},
         expected_user_resolves=[],
         expected_tools=[Tool1, Tool3],
     )
 
     # If none are specifically requested, return all.
     assert_chosen(
-        [], expected_user_resolves=["u1", "u2", "u3"], expected_tools=[Tool1, Tool2, Tool3]
+        set(),
+        expected_user_resolves=[Lang1Requested(["u1", "u2"]), Lang2Requested(["u3"])],
+        expected_tools=[Tool1, Tool2, Tool3],
     )
 
     with pytest.raises(UnrecognizedResolveNamesError):
-        assert_chosen(["fake"], expected_user_resolves=[], expected_tools=[])
+        assert_chosen({"fake"}, expected_user_resolves=[], expected_tools=[])
 
+    # TODO: Add ambiguity checks.
     # Error if same resolve name used for tool lockfiles and user lockfiles.
-    class AmbiguousTool(ToolLockfileSentinel):
-        options_scope = "ambiguous"
-
-    with pytest.raises(AmbiguousResolveNamesError):
-        determine_resolves_to_generate(
-            {"ambiguous": "lockfile.txt"}, [AmbiguousTool], ["ambiguous"]
-        )
+    # class AmbiguousTool(ToolLockfileSentinel):
+    #     options_scope = "ambiguous"
+    #
+    # with pytest.raises(AmbiguousResolveNamesError):
+    #     determine_resolves_to_generate(
+    #         {"ambiguous": "lockfile.txt"}, [AmbiguousTool], ["ambiguous"]
+    #     )
 
 
 def test_filter_tool_lockfile_requests() -> None: