Plumb entire lockfile/internal only code through `create_pex_from_tar…

…gets` (cherry-pick of #18622) (#18636)

Previously if `_determine_requirements_for_pex_from_targets` returns a
`PexRequest` we'd short-circuit the rest of of the code resulting in
`main` not being set on the PEX that we run (most likely in addition to
other bugs like local dists not existing either).

Refactored so that in this very specific case, we'd still make it
through the rest of the code, leveraging `pex_path` for the repo PEX.

Fixes #18552

Co-authored-by: Joshua Cannon <joshdcannon@gmail.com>

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

@@ -26,6 +26,7 @@
     CompletePlatforms,
     OptionalPex,
     OptionalPexRequest,
+    Pex,
     PexPlatforms,
     PexRequest,
 )
@@ -373,9 +374,9 @@ class _ConstraintsRepositoryPexRequest:
 @rule_helper
 async def _determine_requirements_for_pex_from_targets(
     request: PexFromTargetsRequest, python_setup: PythonSetup
-) -> PexRequirements | PexRequest:
+) -> tuple[PexRequirements | EntireLockfile, Iterable[Pex]]:
     if not request.include_requirements:
-        return PexRequirements()
+        return PexRequirements(), ()
 
     requirements = await Get(PexRequirements, _PexRequirementsRequest(request.addresses))
     pex_native_subsetting_supported = False
@@ -409,13 +410,13 @@ async def _determine_requirements_for_pex_from_targets(
 
     if not should_request_repository_pex:
         if not pex_native_subsetting_supported:
-            return requirements
+            return requirements, ()
 
         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=loaded_lockfile), ()
 
     # Else, request the repository PEX and possibly subset it.
     repository_pex_request = await Get(
@@ -440,22 +441,23 @@ async def _determine_requirements_for_pex_from_targets(
                     """
                 )
             )
-        return repository_pex_request.maybe_pex_request
 
     repository_pex = await Get(OptionalPex, OptionalPexRequest, repository_pex_request)
-    return dataclasses.replace(requirements, from_superset=repository_pex.maybe_pex)
+    if should_return_entire_lockfile:
+        assert repository_pex_request.maybe_pex_request is not None
+        assert repository_pex.maybe_pex is not None
+        return repository_pex_request.maybe_pex_request.requirements, [repository_pex.maybe_pex]
+
+    return dataclasses.replace(requirements, from_superset=repository_pex.maybe_pex), ()
 
 
 @rule(level=LogLevel.DEBUG)
 async def create_pex_from_targets(
     request: PexFromTargetsRequest, python_setup: PythonSetup
 ) -> PexRequest:
-    requirements_or_pex_request = await _determine_requirements_for_pex_from_targets(
+    requirements, additional_pexes = await _determine_requirements_for_pex_from_targets(
         request, python_setup
     )
-    if isinstance(requirements_or_pex_request, PexRequest):
-        return requirements_or_pex_request
-    requirements = requirements_or_pex_request
 
     interpreter_constraints = await Get(
         InterpreterConstraints,
@@ -519,6 +521,7 @@ async def create_pex_from_targets(
         additional_inputs=additional_inputs,
         additional_args=additional_args,
         description=description,
+        pex_path=additional_pexes,
     )
 
 

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

@@ -19,6 +19,7 @@
 from pants.backend.python.subsystems.setup import PythonSetup
 from pants.backend.python.subsystems.setuptools import Setuptools
 from pants.backend.python.target_types import (
+    EntryPoint,
     PexLayout,
     PythonRequirementTarget,
     PythonSourcesGeneratorTarget,
@@ -180,7 +181,8 @@ def assert_setup(
         platforms: bool,
         include_requirements: bool = True,
         run_against_entire_lockfile: bool = False,
-        expected: PexRequirements | PexRequest,
+        expected_reqs: PexRequirements = PexRequirements(),
+        expected_pexes: Iterable[Pex] = (),
     ) -> None:
         lockfile_used = mode in (RequirementMode.PEX_LOCKFILE, RequirementMode.NON_PEX_LOCKFILE)
         requirement_constraints_used = mode in (
@@ -224,7 +226,7 @@ def assert_setup(
             mock_repository_pex_request = OptionalPexRequest(maybe_pex_request=None)
             mock_repository_pex = OptionalPex(maybe_pex=None)
 
-        requirements_or_pex_request = run_rule_with_mocks(
+        reqs, pexes = run_rule_with_mocks(
             _determine_requirements_for_pex_from_targets,
             rule_args=[pex_from_targets_request, python_setup],
             mock_gets=[
@@ -259,8 +261,8 @@ def assert_setup(
                 ),
             ],
         )
-        if expected:
-            assert requirements_or_pex_request == expected
+        assert expected_reqs == reqs
+        assert expected_pexes == pexes
 
     # If include_requirements is False, no matter what, early return.
     for mode in RequirementMode:
@@ -269,7 +271,7 @@ def assert_setup(
             include_requirements=False,
             internal_only=False,
             platforms=False,
-            expected=PexRequirements(),
+            # Nothing is expected
         )
 
     # Pex lockfiles: usually, return PexRequirements with from_superset as the LoadedLockfile.
@@ -280,28 +282,30 @@ def assert_setup(
             RequirementMode.PEX_LOCKFILE,
             internal_only=internal_only,
             platforms=False,
-            expected=PexRequirements(req_strings, from_superset=loaded_lockfile__pex),
+            expected_reqs=PexRequirements(req_strings, from_superset=loaded_lockfile__pex),
         )
+
     assert_setup(
         RequirementMode.PEX_LOCKFILE,
         internal_only=False,
         platforms=True,
-        expected=PexRequirements(req_strings, from_superset=loaded_lockfile__pex),
+        expected_reqs=PexRequirements(req_strings, from_superset=loaded_lockfile__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_reqs=PexRequirements(req_strings, from_superset=loaded_lockfile__pex),
         )
     assert_setup(
         RequirementMode.PEX_LOCKFILE,
         internal_only=True,
         run_against_entire_lockfile=True,
         platforms=False,
-        expected=repository_pex_request__lockfile,
+        expected_reqs=repository_pex_request__lockfile.requirements,
+        expected_pexes=[repository_pex__lockfile],
     )
 
     # Non-Pex lockfiles: except for when run_against_entire_lockfile is applicable, return
@@ -312,22 +316,24 @@ def assert_setup(
             RequirementMode.NON_PEX_LOCKFILE,
             internal_only=internal_only,
             platforms=False,
-            expected=PexRequirements(
+            expected_reqs=PexRequirements(
                 req_strings, constraints_strings=req_strings, from_superset=repository_pex__lockfile
             ),
         )
     assert_setup(
         RequirementMode.NON_PEX_LOCKFILE,
         internal_only=False,
         platforms=True,
-        expected=PexRequirements(req_strings, constraints_strings=req_strings, from_superset=None),
+        expected_reqs=PexRequirements(
+            req_strings, constraints_strings=req_strings, from_superset=None
+        ),
     )
     assert_setup(
         RequirementMode.NON_PEX_LOCKFILE,
         internal_only=False,
         run_against_entire_lockfile=True,
         platforms=False,
-        expected=PexRequirements(
+        expected_reqs=PexRequirements(
             req_strings, constraints_strings=req_strings, from_superset=repository_pex__lockfile
         ),
     )
@@ -336,14 +342,17 @@ def assert_setup(
         internal_only=False,
         run_against_entire_lockfile=True,
         platforms=True,
-        expected=PexRequirements(req_strings, constraints_strings=req_strings, from_superset=None),
+        expected_reqs=PexRequirements(
+            req_strings, constraints_strings=req_strings, from_superset=None
+        ),
     )
     assert_setup(
         RequirementMode.NON_PEX_LOCKFILE,
         internal_only=True,
         run_against_entire_lockfile=True,
         platforms=False,
-        expected=repository_pex_request__lockfile,
+        expected_reqs=repository_pex_request__lockfile.requirements,
+        expected_pexes=[repository_pex__lockfile],
     )
 
     # Constraints file with resolve_all_constraints: except for when run_against_entire_lockfile
@@ -354,7 +363,7 @@ def assert_setup(
             RequirementMode.CONSTRAINTS_RESOLVE_ALL,
             internal_only=internal_only,
             platforms=False,
-            expected=PexRequirements(
+            expected_reqs=PexRequirements(
                 req_strings,
                 constraints_strings=global_requirement_constraints,
                 from_superset=repository_pex__constraints,
@@ -364,7 +373,7 @@ def assert_setup(
         RequirementMode.CONSTRAINTS_RESOLVE_ALL,
         internal_only=False,
         platforms=True,
-        expected=PexRequirements(
+        expected_reqs=PexRequirements(
             req_strings, constraints_strings=global_requirement_constraints, from_superset=None
         ),
     )
@@ -373,7 +382,7 @@ def assert_setup(
         internal_only=False,
         run_against_entire_lockfile=True,
         platforms=False,
-        expected=PexRequirements(
+        expected_reqs=PexRequirements(
             req_strings,
             constraints_strings=global_requirement_constraints,
             from_superset=repository_pex__constraints,
@@ -384,7 +393,7 @@ def assert_setup(
         internal_only=False,
         run_against_entire_lockfile=True,
         platforms=True,
-        expected=PexRequirements(
+        expected_reqs=PexRequirements(
             req_strings, constraints_strings=global_requirement_constraints, from_superset=None
         ),
     )
@@ -393,7 +402,8 @@ def assert_setup(
         internal_only=True,
         run_against_entire_lockfile=True,
         platforms=False,
-        expected=repository_pex_request__constraints,
+        expected_reqs=repository_pex_request__constraints.requirements,
+        expected_pexes=[repository_pex__constraints],
     )
 
     # Constraints file without resolve_all_constraints: always PexRequirements with
@@ -403,7 +413,7 @@ def assert_setup(
             RequirementMode.CONSTRAINTS_NO_RESOLVE_ALL,
             internal_only=internal_only,
             platforms=platforms,
-            expected=PexRequirements(
+            expected_reqs=PexRequirements(
                 req_strings, constraints_strings=global_requirement_constraints
             ),
         )
@@ -412,7 +422,7 @@ def assert_setup(
             RequirementMode.CONSTRAINTS_NO_RESOLVE_ALL,
             internal_only=False,
             platforms=platforms,
-            expected=PexRequirements(
+            expected_reqs=PexRequirements(
                 req_strings, constraints_strings=global_requirement_constraints
             ),
         )
@@ -423,13 +433,13 @@ def assert_setup(
             RequirementMode.NO_LOCKS,
             internal_only=internal_only,
             platforms=False,
-            expected=PexRequirements(req_strings),
+            expected_reqs=PexRequirements(req_strings),
         )
     assert_setup(
         RequirementMode.NO_LOCKS,
         internal_only=False,
         platforms=True,
-        expected=PexRequirements(req_strings),
+        expected_reqs=PexRequirements(req_strings),
     )
 
 
@@ -834,10 +844,12 @@ def test_lockfile_requirements_selection(
         [Address("", target_name="lib")],
         output_filename="demo.pex",
         internal_only=internal_only,
+        main=EntryPoint("a"),
     )
     rule_runner.set_options(options, env_inherit={"PATH"})
     result = rule_runner.request(PexRequest, [request])
     assert result.layout == (PexLayout.PACKED if internal_only else PexLayout.ZIPAPP)
+    assert result.main == EntryPoint("a")
 
     if run_against_entire_lockfile and internal_only:
         # With `run_against_entire_lockfile`, all internal requests result in the full set