From d5df1b16bcb46beef133d4fc64dd2cbbe0b34367 Mon Sep 17 00:00:00 2001 From: Eric Arellano <14852634+Eric-Arellano@users.noreply.github.com> Date: Wed, 10 Aug 2022 11:03:18 -0500 Subject: [PATCH] Refactor `build_pex` rule with `@rule_helper` (#16465) This allows us to leverage early returns & inlining variables. It also allows creating comprehensive unit tests. [ci skip-rust] --- .../pants/backend/python/util_rules/pex.py | 244 +++++++++++------- .../backend/python/util_rules/pex_test.py | 189 +++++++++++++- 2 files changed, 334 insertions(+), 99 deletions(-) diff --git a/src/python/pants/backend/python/util_rules/pex.py b/src/python/pants/backend/python/util_rules/pex.py index 1f0c18db16b..13d9d56aa85 100644 --- a/src/python/pants/backend/python/util_rules/pex.py +++ b/src/python/pants/backend/python/util_rules/pex.py @@ -58,7 +58,7 @@ from pants.engine.internals.selectors import MultiGet from pants.engine.platform import Platform 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, rule, rule_helper from pants.engine.target import HydratedSources, HydrateSourcesRequest, SourcesField, Targets from pants.util.frozendict import FrozenDict from pants.util.logging import LogLevel @@ -354,25 +354,14 @@ def create_pex(self) -> Pex: return Pex(digest=self.digest, name=self.pex_filename, python=self.python) -@rule(level=LogLevel.DEBUG) -async def build_pex( - request: PexRequest, - python_setup: PythonSetup, - python_repos: PythonRepos, - platform: Platform, - pex_runtime_env: PexRuntimeEnvironment, -) -> BuildPexResult: - """Returns a PEX with the given settings.""" - argv = [ - "--output-file", - request.output_filename, - "--no-emit-warnings", - *python_setup.manylinux_pex_args, - *request.additional_args, - ] +@dataclass +class _BuildPexPythonSetup: + python: PythonExecutable | None + argv: list[str] - python: PythonExecutable | None = None +@rule_helper +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 # flags are mutually exclusive. See https://github.com/pantsbuild/pex/issues/957. @@ -380,9 +369,15 @@ async def build_pex( # Note that this means that this is not an internal-only pex. # TODO(#9560): consider validating that these platforms are valid with the interpreter # constraints. - argv.extend(request.platforms.generate_pex_arg_list()) - argv.extend(request.complete_platforms.generate_pex_arg_list()) - elif request.python: + return _BuildPexPythonSetup( + None, + [ + *request.platforms.generate_pex_arg_list(), + *request.complete_platforms.generate_pex_arg_list(), + ], + ) + + if request.python: python = request.python elif request.internal_only: # NB: If it's an internal_only PEX, we do our own lookup of the interpreter based on the @@ -394,30 +389,25 @@ async def build_pex( else: # `--interpreter-constraint` options are mutually exclusive with the `--python` option, # so we only specify them if we have not already located a concrete Python. - argv.extend(request.interpreter_constraints.generate_pex_arg_list()) + return _BuildPexPythonSetup(None, request.interpreter_constraints.generate_pex_arg_list()) - if python: - argv.extend(["--python", python.path]) + return _BuildPexPythonSetup(python, ["--python", python.path]) - if request.main is not None: - argv.extend(request.main.iter_pex_args()) - # TODO(John Sirois): Right now any request requirements will shadow corresponding pex path - # requirements, which could lead to problems. Support shading python binaries. - # See: https://github.com/pantsbuild/pants/issues/9206 - if request.pex_path: - argv.extend(["--pex-path", ":".join(pex.name for pex in request.pex_path)]) +@dataclass +class _BuildPexRequirementsSetup: + digests: list[Digest] + argv: list[str] + concurrency_available: int - source_dir_name = "source_files" - 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_digests = [] - pex_lock_resolver_args = [*python_repos.pex_args] +@rule_helper +async def _setup_pex_requirements( + request: PexRequest, python_repos: PythonRepos, 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"] + if isinstance(request.requirements, EntireLockfile): lockfile, resolve_config = await MultiGet( Get(LoadedLockfile, LoadedLockfileRequest(request.requirements.lockfile)), @@ -426,15 +416,13 @@ async def build_pex( ResolvePexConfigRequest(request.requirements.lockfile.resolve_name), ), ) - concurrency_available = lockfile.requirement_estimate - requirements_digests.append(lockfile.lockfile_digest) - if lockfile.is_pex_native: - argv.extend(["--lock", lockfile.lockfile_path]) - argv.extend(pex_lock_resolver_args) - else: + argv = ( + ["--lock", lockfile.lockfile_path, *pex_lock_resolver_args] + if lockfile.is_pex_native + else # We use pip to resolve a requirements.txt pseudo-lockfile, possibly with hashes. - argv.extend(["--requirement", lockfile.lockfile_path, "--no-transitive"]) - argv.extend(pip_resolver_args) + ["--requirement", lockfile.lockfile_path, "--no-transitive", *pip_resolver_args] + ) if lockfile.metadata and request.requirements.complete_req_strings: validate_metadata( lockfile.metadata, @@ -444,54 +432,111 @@ async def build_pex( python_setup, constraints_file_path_and_hash=resolve_config.constraints_file_path_and_hash, ) - else: - # 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(request.requirements.req_strings) - argv.extend(request.requirements.req_strings) - - if isinstance(request.requirements.from_superset, Pex): - repository_pex = request.requirements.from_superset - argv.extend(["--pex-repository", repository_pex.name]) - requirements_digests.append(repository_pex.digest) - elif isinstance(request.requirements.from_superset, LoadedLockfile): - loaded_lockfile = request.requirements.from_superset - resolve_config = await Get( - ResolvePexConfig, - ResolvePexConfigRequest(loaded_lockfile.original_lockfile.resolve_name), + + return _BuildPexRequirementsSetup( + [lockfile.lockfile_digest], argv, lockfile.requirement_estimate + ) + + # 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(request.requirements.req_strings) + + if isinstance(request.requirements.from_superset, Pex): + repository_pex = request.requirements.from_superset + return _BuildPexRequirementsSetup( + [repository_pex.digest], + [*request.requirements.req_strings, "--pex-repository", repository_pex.name], + concurrency_available, + ) + + if isinstance(request.requirements.from_superset, LoadedLockfile): + 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) + + if loaded_lockfile.metadata: + validate_metadata( + loaded_lockfile.metadata, + request.interpreter_constraints, + loaded_lockfile.original_lockfile, + request.requirements.req_strings, + python_setup, + constraints_file_path_and_hash=resolve_config.constraints_file_path_and_hash, ) - # NB: This is also validated in the constructor. - assert loaded_lockfile.is_pex_native - if request.requirements.req_strings: - requirements_digests.append(loaded_lockfile.lockfile_digest) - argv.extend(["--lock", loaded_lockfile.lockfile_path]) - argv.extend(pex_lock_resolver_args) - - if loaded_lockfile.metadata: - validate_metadata( - loaded_lockfile.metadata, - request.interpreter_constraints, - loaded_lockfile.original_lockfile, - request.requirements.req_strings, - python_setup, - constraints_file_path_and_hash=resolve_config.constraints_file_path_and_hash, - ) - else: - assert request.requirements.from_superset is None - - # We use pip to perform a normal resolve. - argv.extend(pip_resolver_args) - if request.requirements.constraints_strings: - constraints_file = "__constraints.txt" - constraints_content = "\n".join(request.requirements.constraints_strings) - requirements_digests.append( - await Get( - Digest, - CreateDigest([FileContent(constraints_file, constraints_content.encode())]), - ) - ) - argv.extend(["--constraints", constraints_file]) + + return _BuildPexRequirementsSetup( + [loaded_lockfile.lockfile_digest], + [ + *request.requirements.req_strings, + "--lock", + loaded_lockfile.lockfile_path, + *pex_lock_resolver_args, + ], + concurrency_available, + ) + + # We use pip to perform a normal resolve. + assert request.requirements.from_superset is None + digests = [] + argv = [*request.requirements.req_strings, *pip_resolver_args] + if request.requirements.constraints_strings: + constraints_file = "__constraints.txt" + constraints_content = "\n".join(request.requirements.constraints_strings) + digests.append( + await Get( + Digest, + CreateDigest([FileContent(constraints_file, constraints_content.encode())]), + ) + ) + argv.extend(["--constraints", constraints_file]) + return _BuildPexRequirementsSetup(digests, argv, concurrency_available=concurrency_available) + + +@rule(level=LogLevel.DEBUG) +async def build_pex( + request: PexRequest, + python_setup: PythonSetup, + python_repos: PythonRepos, + platform: Platform, + pex_runtime_env: PexRuntimeEnvironment, +) -> BuildPexResult: + """Returns a PEX with the given settings.""" + argv = [ + "--output-file", + request.output_filename, + "--no-emit-warnings", + *python_setup.manylinux_pex_args, + *request.additional_args, + ] + + pex_python_setup = await _determine_pex_python_and_platforms(request) + argv.extend(pex_python_setup.argv) + + if request.main is not None: + argv.extend(request.main.iter_pex_args()) + + # TODO(John Sirois): Right now any request requirements will shadow corresponding pex path + # requirements, which could lead to problems. Support shading python binaries. + # See: https://github.com/pantsbuild/pants/issues/9206 + if request.pex_path: + argv.extend(["--pex-path", ":".join(pex.name for pex in request.pex_path)]) + + source_dir_name = "source_files" + 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_repos, python_setup) + argv.extend(requirements_setup.argv) merged_digest = await Get( Digest, @@ -500,7 +545,7 @@ async def build_pex( request.complete_platforms.digest, sources_digest_as_subdir, request.additional_inputs, - *requirements_digests, + *requirements_setup.digests, *(pex.digest for pex in request.pex_path), ) ), @@ -517,14 +562,14 @@ async def build_pex( process = await Get( Process, PexCliProcess( - python=python, + python=pex_python_setup.python, subcommand=(), extra_args=argv, additional_input_digest=merged_digest, description=_build_pex_description(request), output_files=output_files, output_directories=output_directories, - concurrency_available=concurrency_available, + concurrency_available=requirements_setup.concurrency_available, ), ) @@ -549,7 +594,10 @@ async def build_pex( ) return BuildPexResult( - result=result, pex_filename=request.output_filename, digest=digest, python=python + result=result, + pex_filename=request.output_filename, + digest=digest, + python=pex_python_setup.python, ) 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 dd150be8038..0b6eecada60 100644 --- a/src/python/pants/backend/python/util_rules/pex_test.py +++ b/src/python/pants/backend/python/util_rules/pex_test.py @@ -14,11 +14,14 @@ 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 from pants.backend.python.util_rules.interpreter_constraints import InterpreterConstraints from pants.backend.python.util_rules.lockfile_metadata import PythonLockfileMetadata from pants.backend.python.util_rules.pex import ( + CompletePlatforms, Pex, PexDistributionInfo, PexPlatforms, @@ -28,13 +31,22 @@ VenvPex, VenvPexProcess, _build_pex_description, + _BuildPexPythonSetup, + _BuildPexRequirementsSetup, + _determine_pex_python_and_platforms, + _setup_pex_requirements, ) from pants.backend.python.util_rules.pex import rules as pex_rules +from pants.backend.python.util_rules.pex_environment import PythonExecutable from pants.backend.python.util_rules.pex_requirements import ( EntireLockfile, + LoadedLockfile, + LoadedLockfileRequest, Lockfile, LockfileContent, PexRequirements, + ResolvePexConfig, + ResolvePexConfigRequest, ) from pants.backend.python.util_rules.pex_test_utils import ( create_pex_and_get_all_data, @@ -45,7 +57,14 @@ from pants.engine.fs import EMPTY_DIGEST, CreateDigest, Digest, Directory, FileContent from pants.engine.process import Process, ProcessCacheScope, ProcessResult from pants.option.global_options import GlobalOptions -from pants.testutil.rule_runner import QueryRule, RuleRunner, engine_error +from pants.testutil.option_util import create_subsystem +from pants.testutil.rule_runner import ( + MockGet, + QueryRule, + RuleRunner, + engine_error, + run_rule_with_mocks, +) from pants.util.dirutil import safe_rmtree @@ -482,6 +501,174 @@ def test_venv_pex_resolve_info(rule_runner: RuleRunner, pex_type: type[Pex | Ven assert dists[4].project_name == "urllib3" +def test_determine_pex_python_and_platforms() -> None: + hardcoded_python = PythonExecutable("hardcoded/python") + discovered_python = PythonExecutable("discovered/python") + ics = InterpreterConstraints(["==3.7"]) + + def assert_setup( + *, + input_python: PythonExecutable | None = None, + platforms: PexPlatforms = PexPlatforms(), + complete_platforms: CompletePlatforms = CompletePlatforms(), + interpreter_constraints: InterpreterConstraints = InterpreterConstraints(), + internal_only: bool = False, + expected: _BuildPexPythonSetup, + ) -> None: + request = PexRequest( + output_filename="foo.pex", + internal_only=internal_only, + python=input_python, + platforms=platforms, + complete_platforms=complete_platforms, + interpreter_constraints=interpreter_constraints, + ) + result = run_rule_with_mocks( + _determine_pex_python_and_platforms, + rule_args=[request], + mock_gets=[ + MockGet(PythonExecutable, InterpreterConstraints, lambda _: discovered_python) + ], + ) + assert result == expected + + assert_setup(expected=_BuildPexPythonSetup(None, [])) + assert_setup( + interpreter_constraints=ics, + expected=_BuildPexPythonSetup(None, ["--interpreter-constraint", "CPython==3.7"]), + ) + assert_setup( + internal_only=True, + interpreter_constraints=ics, + expected=_BuildPexPythonSetup(discovered_python, ["--python", discovered_python.path]), + ) + assert_setup( + internal_only=True, + input_python=hardcoded_python, + expected=_BuildPexPythonSetup(hardcoded_python, ["--python", hardcoded_python.path]), + ) + assert_setup( + platforms=PexPlatforms(["plat"]), + interpreter_constraints=ics, + expected=_BuildPexPythonSetup(None, ["--platform", "plat"]), + ) + assert_setup( + complete_platforms=CompletePlatforms(["plat"]), + interpreter_constraints=ics, + expected=_BuildPexPythonSetup(None, ["--complete-platform", "plat"]), + ) + + +def test_setup_pex_requirements() -> None: + rule_runner = RuleRunner() + + reqs = ("req1", "req2") + + constraints_content = "constraint" + constraints_digest = rule_runner.make_snapshot( + {"__constraints.txt": constraints_content} + ).digest + + lockfile_path = "foo.lock" + lockfile_digest = rule_runner.make_snapshot_of_empty_files([lockfile_path]).digest + lockfile_obj = Lockfile( + lockfile_path, file_path_description_of_origin="foo", resolve_name="resolve" + ) + + def create_loaded_lockfile(is_pex_lock: bool) -> LoadedLockfile: + return LoadedLockfile( + lockfile_digest, + lockfile_path, + metadata=None, + requirement_estimate=2, + is_pex_native=is_pex_lock, + constraints_strings=None, + original_lockfile=lockfile_obj, + ) + + def assert_setup( + requirements: PexRequirements | EntireLockfile, + expected: _BuildPexRequirementsSetup, + *, + is_pex_lock: bool = True, + ) -> None: + request = PexRequest( + output_filename="foo.pex", + internal_only=True, + requirements=requirements, + ) + result = run_rule_with_mocks( + _setup_pex_requirements, + rule_args=[ + request, + create_subsystem(PythonRepos, indexes=[], repos=[]), + create_subsystem(PythonSetup), + ], + mock_gets=[ + MockGet( + LoadedLockfile, + LoadedLockfileRequest, + lambda _: create_loaded_lockfile(is_pex_lock), + ), + MockGet( + ResolvePexConfig, + ResolvePexConfigRequest, + lambda _: ResolvePexConfig( + constraints_file=None, + ), + ), + MockGet(Digest, CreateDigest, lambda _: constraints_digest), + ], + ) + assert result == expected + + pip_args = ["--no-pypi", "--resolver-version", "pip-2020-resolver"] + pex_args = ["--no-pypi"] + + # Normal resolves. + assert_setup(PexRequirements(reqs), _BuildPexRequirementsSetup([], [*reqs, *pip_args], 2)) + assert_setup( + PexRequirements(reqs, constraints_strings=["constraint"]), + _BuildPexRequirementsSetup( + [constraints_digest], [*reqs, *pip_args, "--constraints", "__constraints.txt"], 2 + ), + ) + + # Pex lockfile. + assert_setup( + EntireLockfile(lockfile_obj, complete_req_strings=reqs), + _BuildPexRequirementsSetup([lockfile_digest], ["--lock", lockfile_path, *pex_args], 2), + ) + + # Non-Pex lockfile. + assert_setup( + EntireLockfile(lockfile_obj, complete_req_strings=reqs), + _BuildPexRequirementsSetup( + [lockfile_digest], ["--requirement", lockfile_path, "--no-transitive", *pip_args], 2 + ), + is_pex_lock=False, + ) + + # Subset of Pex lockfile. + assert_setup( + PexRequirements(["req1"], from_superset=create_loaded_lockfile(is_pex_lock=True)), + _BuildPexRequirementsSetup( + [lockfile_digest], ["req1", "--lock", lockfile_path, *pex_args], 1 + ), + ) + + # Subset of repository Pex. + repository_pex_digest = rule_runner.make_snapshot_of_empty_files(["foo.pex"]).digest + assert_setup( + PexRequirements( + ["req1"], from_superset=Pex(digest=repository_pex_digest, name="foo.pex", python=None) + ), + _BuildPexRequirementsSetup( + [repository_pex_digest], ["req1", "--pex-repository", "foo.pex"], 1 + ), + ) + + def test_build_pex_description() -> None: def assert_description( requirements: PexRequirements | EntireLockfile,