diff --git a/src/python/pants/backend/python/util_rules/pex.py b/src/python/pants/backend/python/util_rules/pex.py index f69af54608b..0a32538e9b2 100644 --- a/src/python/pants/backend/python/util_rules/pex.py +++ b/src/python/pants/backend/python/util_rules/pex.py @@ -63,8 +63,9 @@ from pants.engine.fs import EMPTY_DIGEST, AddPrefix, CreateDigest, Digest, FileContent, MergeDigests from pants.engine.internals.native_engine import Snapshot from pants.engine.internals.selectors import MultiGet +from pants.engine.intrinsics import add_prefix_request_to_digest from pants.engine.process import Process, ProcessCacheScope, ProcessResult -from pants.engine.rules import Get, collect_rules, rule +from pants.engine.rules import Get, collect_rules, concurrently, implicitly, rule from pants.engine.target import ( HydratedSources, HydrateSourcesRequest, @@ -418,6 +419,7 @@ class _BuildPexPythonSetup: argv: list[str] +@rule async def _determine_pex_python_and_platforms(request: PexRequest) -> _BuildPexPythonSetup: # NB: If `--platform` is specified, this signals that the PEX should not be built locally. # `--interpreter-constraint` only makes sense in the context of building locally. These two @@ -509,6 +511,7 @@ async def get_req_strings(pex_reqs: PexRequirements) -> PexRequirementsInfo: return PexRequirementsInfo(tuple(sorted(req_strings)), tuple(sorted(find_links))) +@rule async def _setup_pex_requirements( request: PexRequest, python_setup: PythonSetup ) -> _BuildPexRequirementsSetup: @@ -665,17 +668,42 @@ async def build_pex( ), ) + source_dir_name = "source_files" + + pex_python_setup_req = _determine_pex_python_and_platforms(request) + requirements_setup_req = _setup_pex_requirements(**implicitly({request: PexRequest})) + sources_digest_as_subdir_req = add_prefix_request_to_digest( + AddPrefix(request.sources or EMPTY_DIGEST, source_dir_name) + ) + if isinstance(request.requirements, PexRequirements): + ( + pex_python_setup, + requirements_setup, + sources_digest_as_subdir, + req_info, + ) = await concurrently( + pex_python_setup_req, + requirements_setup_req, + sources_digest_as_subdir_req, + get_req_strings(request.requirements), + ) + req_strings = req_info.req_strings + else: + pex_python_setup, requirements_setup, sources_digest_as_subdir = await concurrently( + pex_python_setup_req, + requirements_setup_req, + sources_digest_as_subdir_req, + ) + req_strings = () + argv = [ "--output-file", request.output_filename, *request.additional_args, ] - pex_python_setup = await _determine_pex_python_and_platforms(request) argv.extend(pex_python_setup.argv) - source_dir_name = "source_files" - if request.main is not None: argv.extend(request.main.iter_pex_args()) if isinstance(request.main, Executable): @@ -705,12 +733,8 @@ async def build_pex( argv.append("--no-pre-install-wheels") argv.append(f"--sources-directory={source_dir_name}") - sources_digest_as_subdir = await Get( - Digest, AddPrefix(request.sources or EMPTY_DIGEST, source_dir_name) - ) # Include any additional arguments and input digests required by the requirements. - requirements_setup = await _setup_pex_requirements(request, python_setup) argv.extend(requirements_setup.argv) merged_digest = await Get( @@ -734,18 +758,13 @@ async def build_pex( else: output_directories = [request.output_filename] - req_strings = ( - (await Get(PexRequirementsInfo, PexRequirements, request.requirements)).req_strings - if isinstance(request.requirements, PexRequirements) - else [] - ) result = await Get( ProcessResult, PexCliProcess( subcommand=(), extra_args=argv, additional_input_digest=merged_digest, - description=await _build_pex_description(request, req_strings, python_setup.resolves), + description=_build_pex_description(request, req_strings, python_setup.resolves), output_files=output_files, output_directories=output_directories, concurrency_available=requirements_setup.concurrency_available, @@ -771,7 +790,7 @@ async def build_pex( ) -async def _build_pex_description( +def _build_pex_description( request: PexRequest, req_strings: Sequence[str], resolve_to_lockfile: Mapping[str, str] ) -> str: if request.description: diff --git a/src/python/pants/backend/python/util_rules/pex_test.py b/src/python/pants/backend/python/util_rules/pex_test.py index 7efa410f667..838dd01d783 100644 --- a/src/python/pants/backend/python/util_rules/pex_test.py +++ b/src/python/pants/backend/python/util_rules/pex_test.py @@ -826,16 +826,12 @@ def assert_description( requirements=requirements, description=description, ) - req_strings = ( - requirements.req_strings_or_addrs if isinstance(requirements, PexRequirements) else [] - ) - assert ( - run_rule_with_mocks( - _build_pex_description, - rule_args=[request, req_strings, {}], - ) - == expected - ) + req_strings = [] + if isinstance(requirements, PexRequirements): + for s in requirements.req_strings_or_addrs: + assert isinstance(s, str) + req_strings.append(s) + assert _build_pex_description(request, req_strings, {}) == expected repo_pex = Pex(EMPTY_DIGEST, "repo.pex", None)