Plumb the hidden _find_links through non-resolve Pex invocations (p…

…antsbuild#20085)

This change fixes pantsbuild#20068 by
adding `find_links` to `ReqStrings`, renaming it `PexRequirementsInfo`,
then plumbing the populated `find_links` through to the relevant
invocation. Namely, the non-subset non-lockfile
`_BuildPexRequirementsSetup` construction.

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

@@ -28,7 +28,13 @@
     PythonLockfileMetadataV2,
     PythonLockfileMetadataV3,
 )
-from pants.backend.python.util_rules.pex import Pex, PexRequest, ReqStrings, VenvPex, VenvPexProcess
+from pants.backend.python.util_rules.pex import (
+    Pex,
+    PexRequest,
+    PexRequirementsInfo,
+    VenvPex,
+    VenvPexProcess,
+)
 from pants.backend.python.util_rules.pex_from_targets import RequirementsPexRequest
 from pants.backend.python.util_rules.pex_requirements import PexRequirements
 from pants.backend.python.util_rules.python_sources import (
@@ -219,7 +225,9 @@ def _count_pytest_tests(contents: DigestContents) -> int:
 async def validate_pytest_cov_included(_pytest: PyTest):
     if _pytest.requirements:
         # We'll only be using this subset of the lockfile.
-        req_strings = (await Get(ReqStrings, PexRequirements(_pytest.requirements))).req_strings
+        req_strings = (
+            await Get(PexRequirementsInfo, PexRequirements(_pytest.requirements))
+        ).req_strings
         requirements = {PipRequirement.parse(req_string) for req_string in req_strings}
     else:
         # We'll be using the entire lockfile.

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

@@ -12,7 +12,7 @@
 from pants.backend.python.subsystems.pytest import PyTest
 from pants.backend.python.util_rules.interpreter_constraints import InterpreterConstraints
 from pants.backend.python.util_rules.lockfile_metadata import PythonLockfileMetadataV3
-from pants.backend.python.util_rules.pex import ReqStrings
+from pants.backend.python.util_rules.pex import PexRequirementsInfo
 from pants.backend.python.util_rules.pex_requirements import (
     LoadedLockfile,
     LoadedLockfileRequest,
@@ -109,7 +109,11 @@ def validate(reqs: list[str]) -> None:
             validate_pytest_cov_included,
             rule_args=[tool],
             mock_gets=[
-                MockGet(ReqStrings, (PexRequirements,), lambda x: ReqStrings(tuple(reqs))),
+                MockGet(
+                    PexRequirementsInfo,
+                    (PexRequirements,),
+                    lambda x: PexRequirementsInfo(tuple(reqs), ()),
+                ),
                 MockGet(Lockfile, (Resolve,), lambda x: lockfile),
                 MockGet(
                     LoadedLockfile,

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

@@ -24,7 +24,10 @@
     PexLayout,
 )
 from pants.backend.python.target_types import PexPlatformsField as PythonPlatformsField
-from pants.backend.python.target_types import PythonRequirementsField
+from pants.backend.python.target_types import (
+    PythonRequirementFindLinksField,
+    PythonRequirementsField,
+)
 from pants.backend.python.util_rules import pex_cli, pex_requirements
 from pants.backend.python.util_rules.interpreter_constraints import InterpreterConstraints
 from pants.backend.python.util_rules.pex_cli import PexCliProcess, PexPEX
@@ -463,15 +466,17 @@ class _BuildPexRequirementsSetup:
 
 
 @dataclass(frozen=True)
-class ReqStrings:
+class PexRequirementsInfo:
     req_strings: tuple[str, ...]
+    find_links: tuple[str, ...]
 
 
 @rule
-async def get_req_strings(pex_reqs: PexRequirements) -> ReqStrings:
+async def get_req_strings(pex_reqs: PexRequirements) -> PexRequirementsInfo:
     addrs: list[Address] = []
     specs: list[str] = []
     req_strings: list[str] = []
+    find_links: set[str] = set()
     for req_str_or_addr in pex_reqs.req_strings_or_addrs:
         if isinstance(req_str_or_addr, Address):
             addrs.append(req_str_or_addr)
@@ -502,7 +507,13 @@ async def get_req_strings(pex_reqs: PexRequirements) -> ReqStrings:
                 if tgt.has_field(PythonRequirementsField)
             )
         )
-    return ReqStrings(tuple(sorted(req_strings)))
+        find_links.update(
+            find_links
+            for tgt in transitive_targets.closure
+            if tgt.has_field(PythonRequirementFindLinksField)
+            for find_links in tgt[PythonRequirementFindLinksField].value or ()
+        )
+    return PexRequirementsInfo(tuple(sorted(req_strings)), tuple(sorted(find_links)))
 
 
 async def _setup_pex_requirements(
@@ -562,18 +573,18 @@ async def _setup_pex_requirements(
         )
 
     assert isinstance(request.requirements, PexRequirements)
-    req_strings = (await Get(ReqStrings, PexRequirements, request.requirements)).req_strings
+    reqs_info = await Get(PexRequirementsInfo, PexRequirements, request.requirements)
 
     # TODO: This is not the best heuristic for available concurrency, since the
     # requirements almost certainly have transitive deps which also need building, but it
     # is better than using something hardcoded.
-    concurrency_available = len(req_strings)
+    concurrency_available = len(reqs_info.req_strings)
 
     if isinstance(request.requirements.from_superset, Pex):
         repository_pex = request.requirements.from_superset
         return _BuildPexRequirementsSetup(
             [repository_pex.digest],
-            [*req_strings, "--pex-repository", repository_pex.name],
+            [*reqs_info.req_strings, "--pex-repository", repository_pex.name],
             concurrency_available,
         )
 
@@ -583,15 +594,15 @@ async def _setup_pex_requirements(
 
         # NB: This is also validated in the constructor.
         assert loaded_lockfile.is_pex_native
-        if not req_strings:
+        if not reqs_info.req_strings:
             return _BuildPexRequirementsSetup([], [], concurrency_available)
 
         if loaded_lockfile.metadata:
             validate_metadata(
                 loaded_lockfile.metadata,
                 request.interpreter_constraints,
                 loaded_lockfile.original_lockfile,
-                consumed_req_strings=req_strings,
+                consumed_req_strings=reqs_info.req_strings,
                 # Don't validate user requirements when subsetting a resolve, as Pex's
                 # validation during the subsetting is far more precise than our naive string
                 # comparison. For example, if a lockfile was generated with `foo==1.2.3`
@@ -605,7 +616,7 @@ async def _setup_pex_requirements(
         return _BuildPexRequirementsSetup(
             [loaded_lockfile.lockfile_digest],
             [
-                *req_strings,
+                *reqs_info.req_strings,
                 "--lock",
                 loaded_lockfile.lockfile_path,
                 *pex_lock_resolver_args,
@@ -615,7 +626,11 @@ async def _setup_pex_requirements(
 
     # We use pip to perform a normal resolve.
     digests = []
-    argv = [*req_strings, *pip_resolver_args]
+    argv = [
+        *reqs_info.req_strings,
+        *pip_resolver_args,
+        *(f"--find-links={find_links}" for find_links in reqs_info.find_links),
+    ]
     if request.requirements.constraints_strings:
         constraints_file = "__constraints.txt"
         constraints_content = "\n".join(request.requirements.constraints_strings)
@@ -713,7 +728,7 @@ async def build_pex(
         output_directories = [request.output_filename]
 
     req_strings = (
-        (await Get(ReqStrings, PexRequirements, request.requirements)).req_strings
+        (await Get(PexRequirementsInfo, PexRequirements, request.requirements)).req_strings
         if isinstance(request.requirements, PexRequirements)
         else []
     )

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

@@ -15,6 +15,8 @@
 
 import pytest
 
+from pants.backend.plugin_development import pants_requirements
+from pants.backend.plugin_development.pants_requirements import PantsRequirementsTargetGenerator
 from pants.backend.python import target_types_rules
 from pants.backend.python.goals import package_pex_binary
 from pants.backend.python.subsystems import setuptools
@@ -35,7 +37,7 @@
     Pex,
     PexPlatforms,
     PexRequest,
-    ReqStrings,
+    PexRequirementsInfo,
 )
 from pants.backend.python.util_rules.pex_from_targets import (
     ChosenPythonResolve,
@@ -73,16 +75,18 @@ def rule_runner() -> PythonRuleRunner:
     return PythonRuleRunner(
         rules=[
             *package_pex_binary.rules(),
+            *pants_requirements.rules(),
             *pex_test_utils.rules(),
             *pex_from_targets.rules(),
             *target_types_rules.rules(),
             QueryRule(PexRequest, (PexFromTargetsRequest,)),
-            QueryRule(ReqStrings, (PexRequirements,)),
+            QueryRule(PexRequirementsInfo, (PexRequirements,)),
             QueryRule(GlobalRequirementConstraints, ()),
             QueryRule(ChosenPythonResolve, [ChosenPythonResolveRequest]),
             *setuptools.rules(),
         ],
         target_types=[
+            PantsRequirementsTargetGenerator,
             PexBinary,
             PythonSourcesGeneratorTarget,
             PythonRequirementTarget,
@@ -655,7 +659,7 @@ def get_pex_request(
     pex_req1 = get_pex_request(constraints1_filename, resolve_all_constraints=False)
     assert isinstance(pex_req1.requirements, PexRequirements)
     assert pex_req1.requirements.constraints_strings == FrozenOrderedSet(constraints1_strings)
-    req_strings_obj1 = rule_runner.request(ReqStrings, (pex_req1.requirements,))
+    req_strings_obj1 = rule_runner.request(PexRequirementsInfo, (pex_req1.requirements,))
     assert req_strings_obj1.req_strings == ("bar==5.5.5", "baz", "foo-bar>=0.1.2", url_req)
 
     pex_req2 = get_pex_request(
@@ -666,7 +670,7 @@ def get_pex_request(
     )
     pex_req2_reqs = pex_req2.requirements
     assert isinstance(pex_req2_reqs, PexRequirements)
-    req_strings_obj2 = rule_runner.request(ReqStrings, (pex_req2_reqs,))
+    req_strings_obj2 = rule_runner.request(PexRequirementsInfo, (pex_req2_reqs,))
     assert req_strings_obj2.req_strings == ("bar==5.5.5", "baz", "foo-bar>=0.1.2", url_req)
     assert isinstance(pex_req2_reqs.from_superset, Pex)
     repository_pex = pex_req2_reqs.from_superset
@@ -690,6 +694,34 @@ def get_pex_request(
     get_pex_request(None, resolve_all_constraints=None)
 
 
+def test_pants_requirement(rule_runner: PythonRuleRunner) -> None:
+    rule_runner.write_files(
+        {
+            "app.py": "",
+            "BUILD": dedent(
+                """
+                pants_requirements(name="pants")
+                python_source(name="app", source="app.py", dependencies=[":pants"])
+                """
+            ),
+        }
+    )
+    args = [
+        "--backend-packages=pants.backend.python",
+        "--backend-packages=pants.backend.plugin_development",
+        "--python-repos-indexes=[]",
+    ]
+    request = PexFromTargetsRequest(
+        [Address("", target_name="app")],
+        output_filename="demo.pex",
+        internal_only=False,
+    )
+    rule_runner.set_options(args, env_inherit={"PATH"})
+    pex_req = rule_runner.request(PexRequest, [request])
+    pex_reqs_info = rule_runner.request(PexRequirementsInfo, [pex_req.requirements])
+    assert pex_reqs_info.find_links == ("https://wheels.pantsbuild.org/simple",)
+
+
 @pytest.mark.parametrize("include_requirements", [False, True])
 def test_exclude_requirements(
     include_requirements: bool, tmp_path: Path, rule_runner: PythonRuleRunner

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

@@ -30,8 +30,8 @@
     PexPlatforms,
     PexProcess,
     PexRequest,
+    PexRequirementsInfo,
     PexResolveInfo,
-    ReqStrings,
     VenvPex,
     VenvPexProcess,
     _build_pex_description,
@@ -703,6 +703,7 @@ def assert_setup(
         expected: _BuildPexRequirementsSetup,
         *,
         is_pex_lock: bool = True,
+        include_find_links: bool = False,
     ) -> None:
         request = PexRequest(
             output_filename="foo.pex",
@@ -724,12 +725,13 @@ def assert_setup(
                     mock=lambda _: create_loaded_lockfile(is_pex_lock),
                 ),
                 MockGet(
-                    output_type=ReqStrings,
+                    output_type=PexRequirementsInfo,
                     input_types=(PexRequirements,),
-                    mock=lambda _: ReqStrings(
+                    mock=lambda _: PexRequirementsInfo(
                         tuple(str(x) for x in requirements.req_strings_or_addrs)
                         if isinstance(requirements, PexRequirements)
-                        else tuple()
+                        else tuple(),
+                        ("imma/link",) if include_find_links else tuple(),
                     ),
                 ),
                 MockGet(
@@ -764,6 +766,11 @@ def assert_setup(
 
     # Normal resolves.
     assert_setup(PexRequirements(reqs), _BuildPexRequirementsSetup([], [*reqs, *pip_args], 2))
+    assert_setup(
+        PexRequirements(reqs),
+        _BuildPexRequirementsSetup([], [*reqs, *pip_args, "--find-links=imma/link"], 2),
+        include_find_links=True,
+    )
     assert_setup(
         PexRequirements(reqs, constraints_strings=["constraint"]),
         _BuildPexRequirementsSetup(