Skip to content

Commit

Permalink
Plumb entire lockfile/internal only code through `create_pex_from_tar…
Browse files Browse the repository at this point in the history
…gets` (pantsbuild#18622)

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 pantsbuild#18552
  • Loading branch information
thejcannon authored Mar 31, 2023
1 parent 5442013 commit 25c7f09
Show file tree
Hide file tree
Showing 2 changed files with 48 additions and 33 deletions.
23 changes: 13 additions & 10 deletions src/python/pants/backend/python/util_rules/pex_from_targets.py
Original file line number Diff line number Diff line change
Expand Up @@ -26,6 +26,7 @@
CompletePlatforms,
OptionalPex,
OptionalPexRequest,
Pex,
PexPlatforms,
PexRequest,
)
Expand Down Expand Up @@ -387,9 +388,9 @@ class _ConstraintsRepositoryPexRequest:

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
Expand Down Expand Up @@ -423,12 +424,12 @@ 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)
)
return dataclasses.replace(requirements, from_superset=Resolve(chosen_resolve.name))
return dataclasses.replace(requirements, from_superset=Resolve(chosen_resolve.name)), ()

# Else, request the repository PEX and possibly subset it.
repository_pex_request = await Get(
Expand All @@ -453,22 +454,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,
Expand Down Expand Up @@ -534,6 +536,7 @@ async def create_pex_from_targets(
additional_inputs=additional_inputs,
additional_args=additional_args,
description=description,
pex_path=additional_pexes,
)


Expand Down
58 changes: 35 additions & 23 deletions src/python/pants/backend/python/util_rules/pex_from_targets_test.py
Original file line number Diff line number Diff line change
Expand Up @@ -20,6 +20,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,
Expand Down Expand Up @@ -188,7 +189,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 (
Expand Down Expand Up @@ -232,7 +234,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=[
Expand Down Expand Up @@ -278,8 +280,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:
Expand All @@ -288,7 +290,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 resolve.
Expand All @@ -299,28 +301,30 @@ def assert_setup(
RequirementMode.PEX_LOCKFILE,
_internal_only=internal_only,
_platforms=False,
expected=PexRequirements(req_strings, from_superset=resolve__pex),
expected_reqs=PexRequirements(req_strings, from_superset=resolve__pex),
)

assert_setup(
RequirementMode.PEX_LOCKFILE,
_internal_only=False,
_platforms=True,
expected=PexRequirements(req_strings, from_superset=resolve__pex),
expected_reqs=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=resolve__pex),
expected_reqs=PexRequirements(req_strings, from_superset=resolve__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
Expand All @@ -331,22 +335,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
),
)
Expand All @@ -355,14 +361,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
Expand All @@ -373,7 +382,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,
Expand All @@ -383,7 +392,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
),
)
Expand All @@ -392,7 +401,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,
Expand All @@ -403,7 +412,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
),
)
Expand All @@ -412,7 +421,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
Expand All @@ -422,7 +432,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
),
)
Expand All @@ -431,7 +441,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
),
)
Expand All @@ -442,13 +452,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),
)


Expand Down Expand Up @@ -852,10 +862,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
Expand Down

0 comments on commit 25c7f09

Please sign in to comment.