Support specifying a named resolve as the superset in a PexRequiremen…

…ts. (pantsbuild#18397)

This named resolve must be a user resolve defined in [python].resolves.

This will allow us to move towards specifying a user resolve as the
source to resolve a tool from. We can't use the existing LoadedLockfile
superset for this, because we don't have access to the PythonSetup
subsystem (which provides the list of resolves) when we construct the
PexRequirements for a tool, so we don't know which lockfile the name
corresponds to.

src/python/pants/backend/python/util_rules/pex.py

@@ -37,11 +37,13 @@
     EntireLockfile,
     LoadedLockfile,
     LoadedLockfileRequest,
+    Lockfile,
 )
 from pants.backend.python.util_rules.pex_requirements import (
     PexRequirements as PexRequirements,  # Explicit re-export.
 )
 from pants.backend.python.util_rules.pex_requirements import (
+    Resolve,
     ResolvePexConfig,
     ResolvePexConfigRequest,
     validate_metadata,
@@ -420,8 +422,8 @@ async def _setup_pex_requirements(
     resolve_name: str | None
     if isinstance(request.requirements, EntireLockfile):
         resolve_name = request.requirements.lockfile.resolve_name
-    elif isinstance(request.requirements.from_superset, LoadedLockfile):
-        resolve_name = request.requirements.from_superset.original_lockfile.resolve_name
+    elif isinstance(request.requirements.from_superset, Resolve):
+        resolve_name = request.requirements.from_superset.name
     else:
         # This implies that, currently, per-resolve options are only configurable for resolves.
         # However, if no resolve is specified, we will still load options that apply to every
@@ -433,26 +435,28 @@ async def _setup_pex_requirements(
     pip_resolver_args = [*resolve_config.pex_args(), "--resolver-version", "pip-2020-resolver"]
 
     if isinstance(request.requirements, EntireLockfile):
-        lockfile = await Get(LoadedLockfile, LoadedLockfileRequest(request.requirements.lockfile))
+        loaded_lockfile = await Get(
+            LoadedLockfile, LoadedLockfileRequest(request.requirements.lockfile)
+        )
         argv = (
-            ["--lock", lockfile.lockfile_path, *pex_lock_resolver_args]
-            if lockfile.is_pex_native
+            ["--lock", loaded_lockfile.lockfile_path, *pex_lock_resolver_args]
+            if loaded_lockfile.is_pex_native
             else
             # We use pip to resolve a requirements.txt pseudo-lockfile, possibly with hashes.
-            ["--requirement", lockfile.lockfile_path, "--no-transitive", *pip_resolver_args]
+            ["--requirement", loaded_lockfile.lockfile_path, "--no-transitive", *pip_resolver_args]
         )
-        if lockfile.metadata and request.requirements.complete_req_strings:
+        if loaded_lockfile.metadata and request.requirements.complete_req_strings:
             validate_metadata(
-                lockfile.metadata,
+                loaded_lockfile.metadata,
                 request.interpreter_constraints,
-                lockfile.original_lockfile,
+                loaded_lockfile.original_lockfile,
                 request.requirements.complete_req_strings,
                 python_setup,
                 resolve_config,
             )
 
         return _BuildPexRequirementsSetup(
-            [lockfile.lockfile_digest], argv, lockfile.requirement_estimate
+            [loaded_lockfile.lockfile_digest], argv, loaded_lockfile.requirement_estimate
         )
 
     # TODO: This is not the best heuristic for available concurrency, since the
@@ -468,8 +472,10 @@ async def _setup_pex_requirements(
             concurrency_available,
         )
 
-    if isinstance(request.requirements.from_superset, LoadedLockfile):
-        loaded_lockfile = request.requirements.from_superset
+    elif isinstance(request.requirements.from_superset, Resolve):
+        lockfile = await Get(Lockfile, Resolve, request.requirements.from_superset)
+        loaded_lockfile = await Get(LoadedLockfile, LoadedLockfileRequest(lockfile))
+
         # NB: This is also validated in the constructor.
         assert loaded_lockfile.is_pex_native
         if not request.requirements.req_strings:
@@ -497,7 +503,6 @@ async def _setup_pex_requirements(
         )
 
     # We use pip to perform a normal resolve.
-    assert request.requirements.from_superset is None
     digests = []
     argv = [*request.requirements.req_strings, *pip_resolver_args]
     if request.requirements.constraints_strings:
@@ -580,7 +585,7 @@ async def build_pex(
             subcommand=(),
             extra_args=argv,
             additional_input_digest=merged_digest,
-            description=_build_pex_description(request),
+            description=_build_pex_description(request, python_setup.resolves),
             output_files=output_files,
             output_directories=output_directories,
             concurrency_available=requirements_setup.concurrency_available,
@@ -608,7 +613,7 @@ async def build_pex(
     )
 
 
-def _build_pex_description(request: PexRequest) -> str:
+def _build_pex_description(request: PexRequest, resolve_to_lockfile: Mapping[str, str]) -> str:
     if request.description:
         return request.description
 
@@ -627,8 +632,11 @@ def _build_pex_description(request: PexRequest) -> str:
                 {', '.join(request.requirements.req_strings)}
                 """
             )
-        elif isinstance(request.requirements.from_superset, LoadedLockfile):
-            lockfile_path = request.requirements.from_superset.lockfile_path
+        elif isinstance(request.requirements.from_superset, Resolve):
+            # At this point we know this is a valid user resolve, so we can assume
+            # it's available in the dict. Nonetheless we use get() so that any weird error
+            # here gives a bad message rather than an outright crash.
+            lockfile_path = resolve_to_lockfile.get(request.requirements.from_superset.name, "")
             return softwrap(
                 f"""
                 Building {pluralize(len(request.requirements.req_strings), 'requirement')}

src/python/pants/backend/python/util_rules/pex_from_targets.py

@@ -36,6 +36,7 @@
     LoadedLockfileRequest,
     Lockfile,
     PexRequirements,
+    Resolve,
 )
 from pants.backend.python.util_rules.python_sources import (
     PythonSourceFiles,
@@ -427,8 +428,7 @@ async def _determine_requirements_for_pex_from_targets(
         chosen_resolve = await Get(
             ChosenPythonResolve, ChosenPythonResolveRequest(request.addresses)
         )
-        loaded_lockfile = await Get(LoadedLockfile, LoadedLockfileRequest(chosen_resolve.lockfile))
-        return dataclasses.replace(requirements, from_superset=loaded_lockfile)
+        return dataclasses.replace(requirements, from_superset=Resolve(chosen_resolve.name))
 
     # Else, request the repository PEX and possibly subset it.
     repository_pex_request = await Get(

src/python/pants/backend/python/util_rules/pex_from_targets_test.py

@@ -47,7 +47,9 @@
     EntireLockfile,
     LoadedLockfile,
     LoadedLockfileRequest,
+    Lockfile,
     PexRequirements,
+    Resolve,
 )
 from pants.backend.python.util_rules.pex_test_utils import get_all_data
 from pants.build_graph.address import Address
@@ -164,8 +166,14 @@ class RequirementMode(Enum):
     req_strings = ["req1", "req2"]
     global_requirement_constraints = ["constraint1", "constraint2"]
 
+    resolve__pex = Resolve("pex")
     loaded_lockfile__pex = Mock(is_pex_native=True, as_constraints_strings=None)
+    chosen_resolve__pex = Mock(lockfile=Mock())
+    chosen_resolve__pex.name = "pex"  # name has special meaning in Mock(), so must set it here.
+    resolve__not_pex = Resolve("not_pex")
     loaded_lockfile__not_pex = Mock(is_pex_native=False, as_constraints_strings=req_strings)
+    chosen_resolve__not_pex = Mock(lockfile=Mock())
+    chosen_resolve__not_pex.name = "not_pex"  # ditto.
 
     repository_pex_request__lockfile = Mock()
     repository_pex_request__constraints = Mock()
@@ -236,7 +244,18 @@ def assert_setup(
                 MockGet(
                     output_type=ChosenPythonResolve,
                     input_types=(ChosenPythonResolveRequest,),
-                    mock=lambda _: Mock(lockfile=Mock()),
+                    mock=lambda _: (
+                        chosen_resolve__pex
+                        if _mode == RequirementMode.PEX_LOCKFILE
+                        else chosen_resolve__not_pex
+                    ),
+                ),
+                MockGet(
+                    output_type=Lockfile,
+                    input_types=(Resolve,),
+                    mock=lambda _: (
+                        resolve__pex if _mode == RequirementMode.PEX_LOCKFILE else resolve__not_pex
+                    ),
                 ),
                 MockGet(
                     output_type=LoadedLockfile,
@@ -272,29 +291,29 @@ def assert_setup(
             expected=PexRequirements(),
         )
 
-    # Pex lockfiles: usually, return PexRequirements with from_superset as the LoadedLockfile.
+    # Pex lockfiles: usually, return PexRequirements with from_superset as the resolve.
     #   Except for when run_against_entire_lockfile is set and it's an internal_only Pex, then
     #   return PexRequest.
     for internal_only in (True, False):
         assert_setup(
             RequirementMode.PEX_LOCKFILE,
             _internal_only=internal_only,
             _platforms=False,
-            expected=PexRequirements(req_strings, from_superset=loaded_lockfile__pex),
+            expected=PexRequirements(req_strings, from_superset=resolve__pex),
         )
     assert_setup(
         RequirementMode.PEX_LOCKFILE,
         _internal_only=False,
         _platforms=True,
-        expected=PexRequirements(req_strings, from_superset=loaded_lockfile__pex),
+        expected=PexRequirements(req_strings, from_superset=resolve__pex),
     )
     for platforms in (True, False):
         assert_setup(
             RequirementMode.PEX_LOCKFILE,
             _internal_only=False,
             run_against_entire_lockfile=True,
             _platforms=platforms,
-            expected=PexRequirements(req_strings, from_superset=loaded_lockfile__pex),
+            expected=PexRequirements(req_strings, from_superset=resolve__pex),
         )
     assert_setup(
         RequirementMode.PEX_LOCKFILE,
@@ -822,6 +841,8 @@ def test_lockfile_requirements_selection(
         # NB: It doesn't matter what the lockfile generator is set to: only what is actually on disk.
         options = [
             "--python-enable-resolves",
+            "--python-default-resolve=myresolve",
+            "--python-resolves={'myresolve':'3rdparty/python/default.lock'}",
         ]
 
     if run_against_entire_lockfile:
@@ -855,5 +876,5 @@ def test_lockfile_requirements_selection(
             assert not get_all_data(rule_runner, result.requirements.from_superset).is_zipapp
         else:
             assert mode == ResolveMode.pex
-            assert isinstance(result.requirements.from_superset, LoadedLockfile)
-            assert result.requirements.from_superset.is_pex_native
+            assert isinstance(result.requirements.from_superset, Resolve)
+            assert result.requirements.from_superset.name == "myresolve"

src/python/pants/backend/python/util_rules/pex_requirements.py

@@ -47,6 +47,14 @@
 logger = logging.getLogger(__name__)
 
 
+@dataclass(frozen=True)
+class Resolve:
+    # A named resolve for a "user lockfile".
+    # Soon to be the only kind of lockfile, as this class will help
+    # get rid of the "tool lockfile" concept.
+    name: str
+
+
 @dataclass(frozen=True)
 class Lockfile:
     url: str
@@ -55,6 +63,18 @@ class Lockfile:
     lockfile_hex_digest: str | None = None
 
 
+@rule
+async def get_lockfile_for_resolve(resolve: Resolve, python_setup: PythonSetup) -> Lockfile:
+    lockfile_path = python_setup.resolves.get(resolve.name)
+    if not lockfile_path:
+        raise ValueError(f"No such resolve: {resolve.name}")
+    return Lockfile(
+        url=lockfile_path,
+        url_description_of_origin=f"the resolve `{resolve.name}`",
+        resolve_name=resolve.name,
+    )
+
+
 @dataclass(frozen=True)
 class LoadedLockfile:
     """A lockfile after loading and header stripping.
@@ -258,14 +278,14 @@ class PexRequirements:
     # If these requirements should be resolved as a subset of either a repository PEX, or a
     # PEX-native lockfile, the superset to use. # NB: Use of a lockfile here asserts that the
     # lockfile is PEX-native, because legacy lockfiles do not support subset resolves.
-    from_superset: Pex | LoadedLockfile | None
+    from_superset: Pex | Resolve | None
 
     def __init__(
         self,
         req_strings: Iterable[str] = (),
         *,
         constraints_strings: Iterable[str] = (),
-        from_superset: Pex | LoadedLockfile | None = None,
+        from_superset: Pex | Resolve | None = None,
     ) -> None:
         """
         :param req_strings: The requirement strings to resolve.
@@ -278,19 +298,6 @@ def __init__(
         )
         object.__setattr__(self, "from_superset", from_superset)
 
-        self.__post_init__()
-
-    def __post_init__(self):
-        if isinstance(self.from_superset, LoadedLockfile) and not self.from_superset.is_pex_native:
-            raise ValueError(
-                softwrap(
-                    f"""
-                    The lockfile {self.from_superset.original_lockfile} was not in PEX's
-                    native format, and so cannot be directly used as a superset.
-                    """
-                )
-            )
-
     @classmethod
     def create_from_requirement_fields(
         cls,

src/python/pants/backend/python/util_rules/pex_test.py

@@ -48,6 +48,7 @@
     LoadedLockfileRequest,
     Lockfile,
     PexRequirements,
+    Resolve,
     ResolvePexConfig,
     ResolvePexConfigRequest,
 )
@@ -709,6 +710,11 @@ def assert_setup(
             _setup_pex_requirements,
             rule_args=[request, create_subsystem(PythonSetup)],
             mock_gets=[
+                MockGet(
+                    output_type=Lockfile,
+                    input_types=(Resolve,),
+                    mock=lambda _: lockfile_obj,
+                ),
                 MockGet(
                     output_type=LoadedLockfile,
                     input_types=(LoadedLockfileRequest,),
@@ -770,7 +776,7 @@ def assert_setup(
 
     # Subset of Pex lockfile.
     assert_setup(
-        PexRequirements(["req1"], from_superset=create_loaded_lockfile(is_pex_lock=True)),
+        PexRequirements(["req1"], from_superset=Resolve("resolve")),
         _BuildPexRequirementsSetup(
             [lockfile_digest], ["req1", "--lock", lockfile_path, *pex_args], 1
         ),
@@ -801,7 +807,7 @@ def assert_description(
             requirements=requirements,
             description=description,
         )
-        assert _build_pex_description(request) == expected
+        assert _build_pex_description(request, {}) == expected
 
     repo_pex = Pex(EMPTY_DIGEST, "repo.pex", None)