[internal] Add experimental_resolve field to pex_binary

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

@@ -4,6 +4,7 @@
 from __future__ import annotations
 
 import logging
+from collections import defaultdict
 from dataclasses import dataclass
 from pathlib import PurePath
 from typing import ClassVar, Iterable, Sequence, cast
@@ -18,13 +19,21 @@
     NO_TOOL_LOCKFILE,
     PythonToolRequirementsBase,
 )
-from pants.backend.python.target_types import EntryPoint
+from pants.backend.python.target_types import (
+    EntryPoint,
+    InterpreterConstraintsField,
+    PythonRequirementsField,
+    PythonResolveField,
+    UnrecognizedResolveNamesError,
+)
 from pants.backend.python.util_rules.interpreter_constraints import InterpreterConstraints
 from pants.backend.python.util_rules.lockfile_metadata import (
     LockfileMetadata,
     calculate_invalidation_digest,
 )
-from pants.backend.python.util_rules.pex import PexRequest, VenvPex, VenvPexProcess
+from pants.backend.python.util_rules.pex import PexRequest, PexRequirements, VenvPex, VenvPexProcess
+from pants.base.specs import AddressSpecs, DescendantAddresses
+from pants.engine.collection import Collection
 from pants.engine.fs import (
     CreateDigest,
     Digest,
@@ -36,6 +45,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 TransitiveTargets, TransitiveTargetsRequest, UnexpandedTargets
 from pants.engine.unions import UnionMembership, union
 from pants.python.python_setup import PythonSetup
 from pants.util.logging import LogLevel
@@ -105,6 +115,7 @@ def custom_command(self) -> str | None:
 @dataclass(frozen=True)
 class PythonLockfile:
     digest: Digest
+    resolve_name: str
     path: str
 
 
@@ -233,7 +244,72 @@ async def generate_lockfile(
     final_lockfile_digest = await Get(
         Digest, CreateDigest([FileContent(req.lockfile_dest, lockfile_with_header)])
     )
-    return PythonLockfile(final_lockfile_digest, req.lockfile_dest)
+    return PythonLockfile(final_lockfile_digest, req.resolve_name, req.lockfile_dest)
+
+
+# --------------------------------------------------------------------------------------
+# User lockfiles
+# --------------------------------------------------------------------------------------
+
+
+class _SpecifiedUserResolves(Collection[str]):
+    pass
+
+
+class _UserLockfileRequests(Collection[PythonLockfileRequest]):
+    pass
+
+
+@rule
+async def setup_user_lockfile_requests(
+    requested: _SpecifiedUserResolves, python_setup: PythonSetup
+) -> _UserLockfileRequests:
+    # First, associate all resolves with their consumers.
+    all_build_targets = await Get(UnexpandedTargets, AddressSpecs([DescendantAddresses("")]))
+    resolves_to_roots = defaultdict(list)
+    for tgt in all_build_targets:
+        if not tgt.has_field(PythonResolveField):
+            continue
+        tgt[PythonResolveField].validate(python_setup)
+        resolve = tgt[PythonResolveField].value
+        if resolve is None:
+            continue
+        resolves_to_roots[resolve].append(tgt.address)
+
+    # 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)
+        )
+
+    requests = (
+        PythonLockfileRequest(
+            requirements.req_strings,
+            interpreter_constraints,
+            resolve_name=resolve,
+            lockfile_dest=python_setup.resolves_to_lockfiles[resolve],
+        )
+        for resolve, requirements, interpreter_constraints in zip(
+            requested, pex_requirements_per_resolve, interpreter_constraints_per_resolve
+        )
+    )
+    return _UserLockfileRequests(requests)
 
 
 # --------------------------------------------------------------------------------------
@@ -252,27 +328,33 @@ async def generate_lockfiles_goal(
     generate_lockfiles_subsystem: GenerateLockfilesSubsystem,
     python_setup: PythonSetup,
 ) -> GenerateLockfilesGoal:
-    _specified_user_lockfiles, specified_tool_sentinels = determine_resolves_to_generate(
+    specified_user_resolves, specified_tool_sentinels = determine_resolves_to_generate(
         python_setup.resolves_to_lockfiles.keys(),
         union_membership[PythonToolLockfileSentinel],
         generate_lockfiles_subsystem.resolve_names,
     )
+
+    specified_user_requests = await Get(
+        _UserLockfileRequests, _SpecifiedUserResolves(specified_user_resolves)
+    )
     specified_tool_requests = await MultiGet(
         Get(PythonLockfileRequest, PythonToolLockfileSentinel, sentinel())
         for sentinel in specified_tool_sentinels
     )
+    applicable_tool_requests = filter_tool_lockfile_requests(
+        specified_tool_requests,
+        resolve_specified=bool(generate_lockfiles_subsystem.resolve_names),
+    )
+
     results = await MultiGet(
         Get(PythonLockfile, PythonLockfileRequest, req)
-        for req in filter_tool_lockfile_requests(
-            specified_tool_requests,
-            resolve_specified=bool(generate_lockfiles_subsystem.resolve_names),
-        )
+        for req in (*specified_user_requests, *applicable_tool_requests)
     )
 
     merged_digest = await Get(Digest, MergeDigests(res.digest for res in results))
     workspace.write_digest(merged_digest)
     for result in results:
-        logger.info(f"Wrote lockfile to {result.path}")
+        logger.info(f"Wrote lockfile for the resolve `{result.resolve_name}` to {result.path}")
 
     return GenerateLockfilesGoal(exit_code=0)
 
@@ -298,24 +380,6 @@ def __init__(self, ambiguous_names: list[str]) -> None:
         )
 
 
-class UnrecognizedResolveNamesError(Exception):
-    def __init__(
-        self, unrecognized_resolve_names: list[str], all_valid_names: Iterable[str]
-    ) -> None:
-        # TODO(#12314): maybe implement "Did you mean?"
-        if len(unrecognized_resolve_names) == 1:
-            unrecognized_str = unrecognized_resolve_names[0]
-            name_description = "name"
-        else:
-            unrecognized_str = str(sorted(unrecognized_resolve_names))
-            name_description = "names"
-        super().__init__(
-            f"Unrecognized resolve {name_description} from the option "
-            f"`--generate-lockfiles-resolve`: {unrecognized_str}\n\n"
-            f"All valid resolve names: {sorted(all_valid_names)}"
-        )
-
-
 def determine_resolves_to_generate(
     all_user_resolves: Iterable[str],
     all_tool_sentinels: Iterable[type[PythonToolLockfileSentinel]],
@@ -356,6 +420,7 @@ def determine_resolves_to_generate(
         raise UnrecognizedResolveNamesError(
             unrecognized_resolve_names,
             {*all_user_resolves, *resolve_names_to_sentinels.keys()},
+            description_of_origin="the option `--generate-lockfiles-resolve`",
         )
 
     return specified_user_resolves, specified_sentinels

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

@@ -9,11 +9,11 @@
     AmbiguousResolveNamesError,
     PythonLockfileRequest,
     PythonToolLockfileSentinel,
-    UnrecognizedResolveNamesError,
     determine_resolves_to_generate,
     filter_tool_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.util_rules.interpreter_constraints import InterpreterConstraints
 from pants.util.ordered_set import FrozenOrderedSet
 
@@ -68,24 +68,6 @@ class AmbiguousTool(PythonToolLockfileSentinel):
         )
 
 
-@pytest.mark.parametrize(
-    "unrecognized,bad_entry_str,name_str",
-    (
-        (["fake"], "fake", "name"),
-        (["fake1", "fake2"], "['fake1', 'fake2']", "names"),
-    ),
-)
-def test_unrecognized_resolve_names_error(
-    unrecognized: list[str], bad_entry_str: str, name_str: str
-) -> None:
-    with pytest.raises(UnrecognizedResolveNamesError) as exc:
-        raise UnrecognizedResolveNamesError(unrecognized, ["valid1", "valid2", "valid3"])
-    assert (
-        f"Unrecognized resolve {name_str} from the option `--generate-lockfiles-resolve`: "
-        f"{bad_entry_str}\n\nAll valid resolve names: ['valid1', 'valid2', 'valid3']"
-    ) in str(exc.value)
-
-
 def test_filter_tool_lockfile_requests() -> None:
     def create_request(name: str, lockfile_dest: str | None = None) -> PythonLockfileRequest:
         return PythonLockfileRequest(

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

@@ -21,6 +21,7 @@
     PexShebangField,
     PexStripEnvField,
     PexZipSafeField,
+    PythonResolveField,
     ResolvedPexEntryPoint,
     ResolvePexEntryPointRequest,
 )
@@ -41,6 +42,7 @@
     targets_with_sources_types,
 )
 from pants.engine.unions import UnionMembership, UnionRule
+from pants.python.python_setup import PythonSetup
 from pants.util.docutil import doc_url
 from pants.util.logging import LogLevel
 
@@ -64,6 +66,7 @@ class PexBinaryFieldSet(PackageFieldSet, RunFieldSet):
     platforms: PythonPlatformsField
     execution_mode: PexExecutionModeField
     include_tools: PexIncludeToolsField
+    resolve: PythonResolveField
 
     @property
     def _execution_mode(self) -> PexExecutionMode:
@@ -98,6 +101,7 @@ def generate_additional_args(self, pex_binary_defaults: PexBinaryDefaults) -> Tu
 async def package_pex_binary(
     field_set: PexBinaryFieldSet,
     pex_binary_defaults: PexBinaryDefaults,
+    python_setup: PythonSetup,
     union_membership: UnionMembership,
 ) -> BuiltPackage:
     resolved_entry_point, transitive_targets = await MultiGet(
@@ -132,6 +136,7 @@ async def package_pex_binary(
             #  https://github.com/pantsbuild/pants/issues/11619
             main=resolved_entry_point.val,
             platforms=PexPlatforms.create_from_platforms_field(field_set.platforms),
+            resolve_and_lockfile=field_set.resolve.resolve_and_lockfile(python_setup),
             output_filename=output_filename,
             additional_args=field_set.generate_additional_args(pex_binary_defaults),
         ),

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

@@ -21,6 +21,7 @@
 from pants.engine.rules import Get, MultiGet, collect_rules, rule
 from pants.engine.target import TransitiveTargets, TransitiveTargetsRequest
 from pants.engine.unions import UnionRule
+from pants.python.python_setup import PythonSetup
 from pants.util.logging import LogLevel
 
 
@@ -29,6 +30,7 @@ async def create_pex_binary_run_request(
     field_set: PexBinaryFieldSet,
     pex_binary_defaults: PexBinaryDefaults,
     pex_env: PexEnvironment,
+    python_setup: PythonSetup,
 ) -> RunRequest:
     entry_point, transitive_targets = await MultiGet(
         Get(
@@ -43,7 +45,11 @@ async def create_pex_binary_run_request(
     requirements_pex_request = await Get(
         PexRequest,
         PexFromTargetsRequest,
-        PexFromTargetsRequest.for_requirements([field_set.address], internal_only=True),
+        PexFromTargetsRequest.for_requirements(
+            [field_set.address],
+            internal_only=True,
+            resolve_and_lockfile=field_set.resolve.resolve_and_lockfile(python_setup),
+        ),
     )
 
     requirements_request = Get(Pex, PexRequest, requirements_pex_request)

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

@@ -86,6 +86,60 @@ def value_or_global_default(self, python_setup: PythonSetup) -> Tuple[str, ...]:
         return python_setup.compatibility_or_constraints(self.value)
 
 
+class UnrecognizedResolveNamesError(Exception):
+    def __init__(
+        self,
+        unrecognized_resolve_names: list[str],
+        all_valid_names: Iterable[str],
+        *,
+        description_of_origin: str,
+    ) -> None:
+        # TODO(#12314): maybe implement "Did you mean?"
+        if len(unrecognized_resolve_names) == 1:
+            unrecognized_str = unrecognized_resolve_names[0]
+            name_description = "name"
+        else:
+            unrecognized_str = str(sorted(unrecognized_resolve_names))
+            name_description = "names"
+        super().__init__(
+            f"Unrecognized resolve {name_description} from {description_of_origin}: "
+            f"{unrecognized_str}\n\nAll valid resolve names: {sorted(all_valid_names)}"
+        )
+
+
+class PythonResolveField(StringField, AsyncFieldMixin):
+    alias = "experimental_resolve"
+    # TODO(#12314): Figure out how to model the default and disabling lockfile, e.g. if we
+    #  hardcode to `default` or let the user set it.
+    help = (
+        "The resolve from `[python-setup].experimental_resolves_to_lockfiles` to use, if any.\n\n"
+        "This field is highly experimental and may change without the normal deprecation policy."
+    )
+
+    def validate(self, python_setup: PythonSetup) -> None:
+        """Check that the resolve name is recognized."""
+        if not self.value:
+            return None
+        if self.value not in python_setup.resolves_to_lockfiles:
+            raise UnrecognizedResolveNamesError(
+                [self.value],
+                python_setup.resolves_to_lockfiles.keys(),
+                description_of_origin=f"the field `{self.alias}` in the target {self.address}",
+            )
+
+    def resolve_and_lockfile(self, python_setup: PythonSetup) -> tuple[str, str] | None:
+        """If configured, return the resolve name with its lockfile.
+
+        Error if the resolve name is invalid.
+        """
+        self.validate(python_setup)
+        return (
+            (self.value, python_setup.resolves_to_lockfiles[self.value])
+            if self.value is not None
+            else None
+        )
+
+
 # -----------------------------------------------------------------------------------------------
 # `pex_binary` target
 # -----------------------------------------------------------------------------------------------
@@ -383,6 +437,7 @@ class PexBinary(Target):
         *COMMON_TARGET_FIELDS,
         OutputPathField,
         InterpreterConstraintsField,
+        PythonResolveField,
         PexBinaryDependencies,
         PexEntryPointField,
         PexPlatformsField,