Refactor how [python-repos] is loaded (pantsbuild#16520)

Prework for:

1) Adding per-resolve options for `[python-repos].{repos,indexes}` and `[python].manylinux` 
2) Wiring up these 3 options to lockfile header metadata pantsbuild#12832

End behavior is the same.

[ci skip-rust]
[ci skip-build-wheels]

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

@@ -156,20 +156,22 @@ def warn_python_repos(option: str) -> None:
 
 @dataclass(frozen=True)
 class _PipArgsAndConstraintsSetup:
+    resolve_config: ResolvePexConfig
     args: tuple[str, ...]
     digest: Digest
-    constraints: FrozenOrderedSet[PipRequirement]
-    only_binary: tuple[str, ...]
-    no_binary: tuple[str, ...]
 
 
 @rule_helper
-async def _setup_pip_args_and_constraints_file(resolve_name: str) -> _PipArgsAndConstraintsSetup:
-    args = []
-    digests = []
+async def _setup_pip_args_and_constraints_file(
+    resolve_name: str, *, generate_with_pex: bool
+) -> _PipArgsAndConstraintsSetup:
     resolve_config = await Get(ResolvePexConfig, ResolvePexConfigRequest(resolve_name))
 
-    if resolve_config.no_binary or resolve_config.only_binary:
+    args = list(resolve_config.indexes_and_find_links_and_manylinux_pex_args())
+    digests = []
+
+    # This feature only works with Pex lockfiles.
+    if generate_with_pex and (resolve_config.no_binary or resolve_config.only_binary):
         pip_args_file = "__pip_args.txt"
         args.extend(["-r", pip_args_file])
         pip_args_file_content = "\n".join(
@@ -181,41 +183,26 @@ async def _setup_pip_args_and_constraints_file(resolve_name: str) -> _PipArgsAnd
         )
         digests.append(pip_args_digest)
 
-    constraints: FrozenOrderedSet[PipRequirement] = FrozenOrderedSet()
-    resolve_config = await Get(ResolvePexConfig, ResolvePexConfigRequest(resolve_name))
-    if resolve_config.constraints_file:
+    # This feature only works with Pex lockfiles.
+    if generate_with_pex and resolve_config.constraints_file:
         args.append(f"--constraints={resolve_config.constraints_file.path}")
         digests.append(resolve_config.constraints_file.digest)
-        constraints = resolve_config.constraints_file.constraints
 
     input_digest = await Get(Digest, MergeDigests(digests))
-    return _PipArgsAndConstraintsSetup(
-        tuple(args),
-        input_digest,
-        constraints,
-        only_binary=resolve_config.only_binary,
-        no_binary=resolve_config.no_binary,
-    )
+    return _PipArgsAndConstraintsSetup(resolve_config, tuple(args), input_digest)
 
 
 @rule(desc="Generate Python lockfile", level=LogLevel.DEBUG)
 async def generate_lockfile(
     req: GeneratePythonLockfile,
     poetry_subsystem: PoetrySubsystem,
     generate_lockfiles_subsystem: GenerateLockfilesSubsystem,
-    python_repos: PythonRepos,
-    python_setup: PythonSetup,
 ) -> GenerateLockfileResult:
-    requirement_constraints: FrozenOrderedSet[PipRequirement] = FrozenOrderedSet()
-    only_binary: tuple[str, ...] = ()
-    no_binary: tuple[str, ...] = ()
+    pip_args_setup = await _setup_pip_args_and_constraints_file(
+        req.resolve_name, generate_with_pex=req.use_pex
+    )
 
     if req.use_pex:
-        pip_args_setup = await _setup_pip_args_and_constraints_file(req.resolve_name)
-        requirement_constraints = pip_args_setup.constraints
-        only_binary = pip_args_setup.only_binary
-        no_binary = pip_args_setup.no_binary
-
         header_delimiter = "//"
         result = await Get(
             ProcessResult,
@@ -245,8 +232,6 @@ async def generate_lockfile(
                     # This makes diffs more readable when lockfiles change.
                     "--indent=2",
                     *pip_args_setup.args,
-                    *python_repos.pex_args,
-                    *python_setup.manylinux_pex_args,
                     *req.interpreter_constraints.generate_pex_arg_list(),
                     *req.requirements,
                 ),
@@ -313,13 +298,17 @@ async def generate_lockfile(
         )
 
     initial_lockfile_digest_contents = await Get(DigestContents, Digest, result.output_digest)
-    # TODO(#12314) Improve error message on `Requirement.parse`
     metadata = PythonLockfileMetadata.new(
         valid_for_interpreter_constraints=req.interpreter_constraints,
+        # TODO(#12314) Improve error message on `Requirement.parse`
         requirements={PipRequirement.parse(i) for i in req.requirements},
-        requirement_constraints=set(requirement_constraints),
-        only_binary=set(only_binary),
-        no_binary=set(no_binary),
+        requirement_constraints=(
+            set(pip_args_setup.resolve_config.constraints_file.constraints)
+            if pip_args_setup.resolve_config.constraints_file
+            else set()
+        ),
+        only_binary=set(pip_args_setup.resolve_config.only_binary),
+        no_binary=set(pip_args_setup.resolve_config.no_binary),
     )
     lockfile_with_header = metadata.add_header_to_lockfile(
         initial_lockfile_digest_contents[0].content,

src/python/pants/backend/python/subsystems/repos.py

@@ -3,8 +3,6 @@
 
 from __future__ import annotations
 
-from typing import Iterator
-
 from pants.option.option_types import StrListOption
 from pants.option.subsystem import Subsystem
 from pants.util.strutil import softwrap
@@ -42,13 +40,3 @@ class PythonRepos(Subsystem):
         ),
         advanced=True,
     )
-
-    @property
-    def pex_args(self) -> Iterator[str]:
-        # NB: In setting `--no-pypi`, we rely on the default value of `--python-repos-indexes`
-        # including PyPI, which will override `--no-pypi` and result in using PyPI in the default
-        # case. Why set `--no-pypi`, then? We need to do this so that
-        # `--python-repos-repos=['custom_url']` will only point to that index and not include PyPI.
-        yield "--no-pypi"
-        yield from (f"--index={index}" for index in self.indexes)
-        yield from (f"--repo={repo}" for repo in self.repos)

src/python/pants/backend/python/subsystems/setup.py

@@ -6,7 +6,7 @@
 import enum
 import logging
 import os
-from typing import Iterable, Iterator, List, Optional, TypeVar, cast
+from typing import Iterable, List, Optional, TypeVar, cast
 
 from pants.core.goals.generate_lockfiles import UnrecognizedResolveNamesError
 from pants.option.option_types import (
@@ -734,14 +734,6 @@ def manylinux(self) -> str | None:
             return None
         return manylinux
 
-    @property
-    def manylinux_pex_args(self) -> Iterator[str]:
-        if self.manylinux:
-            yield "--manylinux"
-            yield self.manylinux
-        else:
-            yield "--no-manylinux"
-
     @property
     def scratch_dir(self):
         return os.path.join(self.options.pants_workdir, *self.options_scope.split("."))

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

@@ -17,7 +17,6 @@
 import packaging.version
 from pkg_resources import Requirement
 
-from pants.backend.python.subsystems.repos import PythonRepos
 from pants.backend.python.subsystems.setup import PythonSetup
 from pants.backend.python.target_types import (
     MainSpecification,
@@ -403,19 +402,29 @@ class _BuildPexRequirementsSetup:
 
 @rule_helper
 async def _setup_pex_requirements(
-    request: PexRequest, python_repos: PythonRepos, python_setup: PythonSetup
+    request: PexRequest, python_setup: PythonSetup
 ) -> _BuildPexRequirementsSetup:
-    pex_lock_resolver_args = list(python_repos.pex_args)
-    pip_resolver_args = [*python_repos.pex_args, "--resolver-version", "pip-2020-resolver"]
+    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
+    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
+        # resolve, like `[python-repos].indexes`.
+        resolve_name = None
+    resolve_config = await Get(ResolvePexConfig, ResolvePexConfigRequest(resolve_name))
+
+    pex_lock_resolver_args = list(resolve_config.indexes_and_find_links_and_manylinux_pex_args())
+    pip_resolver_args = [
+        *resolve_config.indexes_and_find_links_and_manylinux_pex_args(),
+        "--resolver-version",
+        "pip-2020-resolver",
+    ]
 
     if isinstance(request.requirements, EntireLockfile):
-        lockfile, resolve_config = await MultiGet(
-            Get(LoadedLockfile, LoadedLockfileRequest(request.requirements.lockfile)),
-            Get(
-                ResolvePexConfig,
-                ResolvePexConfigRequest(request.requirements.lockfile.resolve_name),
-            ),
-        )
+        lockfile = await Get(LoadedLockfile, LoadedLockfileRequest(request.requirements.lockfile))
         argv = (
             ["--lock", lockfile.lockfile_path, *pex_lock_resolver_args]
             if lockfile.is_pex_native
@@ -454,10 +463,6 @@ async def _setup_pex_requirements(
         loaded_lockfile = request.requirements.from_superset
         # NB: This is also validated in the constructor.
         assert loaded_lockfile.is_pex_native
-        resolve_config = await Get(
-            ResolvePexConfig,
-            ResolvePexConfigRequest(loaded_lockfile.original_lockfile.resolve_name),
-        )
         if not request.requirements.req_strings:
             return _BuildPexRequirementsSetup([], [], concurrency_available)
 
@@ -503,7 +508,6 @@ async def _setup_pex_requirements(
 async def build_pex(
     request: PexRequest,
     python_setup: PythonSetup,
-    python_repos: PythonRepos,
     platform: Platform,
     pex_runtime_env: PexRuntimeEnvironment,
 ) -> BuildPexResult:
@@ -512,7 +516,6 @@ async def build_pex(
         "--output-file",
         request.output_filename,
         "--no-emit-warnings",
-        *python_setup.manylinux_pex_args,
         *request.additional_args,
     ]
 
@@ -535,7 +538,7 @@ async def build_pex(
     )
 
     # Include any additional arguments and input digests required by the requirements.
-    requirements_setup = await _setup_pex_requirements(request, python_repos, python_setup)
+    requirements_setup = await _setup_pex_requirements(request, python_setup)
     argv.extend(requirements_setup.argv)
 
     merged_digest = await Get(

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

@@ -8,6 +8,7 @@
 from typing import TYPE_CHECKING, Iterable, Iterator
 
 from pants.backend.python.pip_requirement import PipRequirement
+from pants.backend.python.subsystems.repos import PythonRepos
 from pants.backend.python.subsystems.setup import InvalidLockfileBehavior, PythonSetup
 from pants.backend.python.target_types import PythonRequirementsField, parse_requirements_file
 from pants.backend.python.util_rules.interpreter_constraints import InterpreterConstraints
@@ -297,25 +298,62 @@ class ResolvePexConstraintsFile:
 class ResolvePexConfig:
     """Configuration from `[python]` that impacts how the resolve is created."""
 
+    indexes: tuple[str, ...]
+    find_links: tuple[str, ...]
+    manylinux: str | None
     constraints_file: ResolvePexConstraintsFile | None
     only_binary: tuple[str, ...]
     no_binary: tuple[str, ...]
 
+    def indexes_and_find_links_and_manylinux_pex_args(self) -> Iterator[str]:
+        # NB: In setting `--no-pypi`, we rely on the default value of `[python-repos].indexes`
+        # including PyPI, which will override `--no-pypi` and result in using PyPI in the default
+        # case. Why set `--no-pypi`, then? We need to do this so that
+        # `[python-repos].indexes = ['custom_url']` will only point to that index and not include
+        # PyPI.
+        yield "--no-pypi"
+        yield from (f"--index={index}" for index in self.indexes)
+        yield from (f"--find-links={repo}" for repo in self.find_links)
+
+        if self.manylinux:
+            yield "--manylinux"
+            yield self.manylinux
+        else:
+            yield "--no-manylinux"
+
 
 @dataclass(frozen=True)
 class ResolvePexConfigRequest(EngineAwareParameter):
-    """Find all configuration from `[python]` that impacts how the resolve is created."""
+    """Find all configuration from `[python]` that impacts how the resolve is created.
 
-    resolve_name: str
+    If `resolve_name` is None, then most per-resolve options will be ignored because there is no way
+    for users to configure them. However, some options like `[python-repos].indexes` will still be
+    loaded.
+    """
+
+    resolve_name: str | None
 
     def debug_hint(self) -> str:
-        return self.resolve_name
+        return self.resolve_name or "<no resolve>"
 
 
 @rule
 async def determine_resolve_pex_config(
-    request: ResolvePexConfigRequest, python_setup: PythonSetup, union_membership: UnionMembership
+    request: ResolvePexConfigRequest,
+    python_setup: PythonSetup,
+    python_repos: PythonRepos,
+    union_membership: UnionMembership,
 ) -> ResolvePexConfig:
+    if request.resolve_name is None:
+        return ResolvePexConfig(
+            indexes=python_repos.indexes,
+            find_links=python_repos.repos,
+            manylinux=python_setup.manylinux,
+            constraints_file=None,
+            no_binary=(),
+            only_binary=(),
+        )
+
     all_python_tool_resolve_names = tuple(
         sentinel.resolve_name
         for sentinel in union_membership.get(GenerateToolLockfileSentinel)
@@ -374,6 +412,9 @@ async def determine_resolve_pex_config(
         )
 
     return ResolvePexConfig(
+        indexes=python_repos.indexes,
+        find_links=python_repos.repos,
+        manylinux=python_setup.manylinux,
         constraints_file=constraints_file,
         no_binary=tuple(no_binary),
         only_binary=tuple(only_binary),

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

@@ -136,7 +136,10 @@ def test_validate_tool_lockfiles(
         req_strings,
         create_python_setup(InvalidLockfileBehavior.warn),
         ResolvePexConfig(
-            ResolvePexConstraintsFile(
+            indexes=(),
+            find_links=(),
+            manylinux=None,
+            constraints_file=ResolvePexConstraintsFile(
                 EMPTY_DIGEST,
                 "c.txt",
                 FrozenOrderedSet(
@@ -242,7 +245,10 @@ def test_validate_user_lockfiles(
         req_strings,
         create_python_setup(InvalidLockfileBehavior.warn),
         ResolvePexConfig(
-            ResolvePexConstraintsFile(
+            indexes=(),
+            find_links=(),
+            manylinux=None,
+            constraints_file=ResolvePexConstraintsFile(
                 EMPTY_DIGEST,
                 "c.txt",
                 FrozenOrderedSet(

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

@@ -14,7 +14,6 @@
 from pkg_resources import Requirement
 
 from pants.backend.python.pip_requirement import PipRequirement
-from pants.backend.python.subsystems.repos import PythonRepos
 from pants.backend.python.subsystems.setup import PythonSetup
 from pants.backend.python.target_types import EntryPoint
 from pants.backend.python.util_rules import pex_test_utils
@@ -599,11 +598,7 @@ def assert_setup(
         )
         result = run_rule_with_mocks(
             _setup_pex_requirements,
-            rule_args=[
-                request,
-                create_subsystem(PythonRepos, indexes=[], repos=[]),
-                create_subsystem(PythonSetup),
-            ],
+            rule_args=[request, create_subsystem(PythonSetup)],
             mock_gets=[
                 MockGet(
                     LoadedLockfile,
@@ -614,6 +609,9 @@ def assert_setup(
                     ResolvePexConfig,
                     ResolvePexConfigRequest,
                     lambda _: ResolvePexConfig(
+                        indexes=("custom-index",),
+                        find_links=("custom-find-links",),
+                        manylinux=None,
                         constraints_file=None,
                         only_binary=(),
                         no_binary=(),
@@ -624,8 +622,13 @@ def assert_setup(
         )
         assert result == expected
 
-    pip_args = ["--no-pypi", "--resolver-version", "pip-2020-resolver"]
-    pex_args = ["--no-pypi"]
+    pex_args = [
+        "--no-pypi",
+        "--index=custom-index",
+        "--find-links=custom-find-links",
+        "--no-manylinux",
+    ]
+    pip_args = [*pex_args, "--resolver-version", "pip-2020-resolver"]
 
     # Normal resolves.
     assert_setup(PexRequirements(reqs), _BuildPexRequirementsSetup([], [*reqs, *pip_args], 2))