24149123771

Committed 08 Apr 2026 05:25PM UTC coverage: 92.925% (+0.02%) from 92.91%

Build # 24149123771

Build Type

Pull #23227

github

Committed by

web-flow

Commit Message

Merge c74227668 into 9036734c9

Pull Request Pull Request #23227: Fix uv PEX builder to use pex3 lock export

Coverage Stats

85 of 86 new or added lines in 2 files covered. (98.84%)

2 existing lines in 2 files now uncovered.

91671 of 98650 relevant lines covered (92.93%)

4.04 hits per line

Source File
Press 'n' to go to next uncovered line, 'b' for previous

99.73

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

# Copyright 2019 Pants project contributors (see CONTRIBUTORS.md).
# Licensed under the Apache License, Version 2.0 (see LICENSE).

from __future__ import annotations

import os.path
import pkgutil
import re
import shutil
import textwrap
import zipfile
from pathlib import Path
from types import SimpleNamespace

import pytest
import requests
from packaging.requirements import Requirement
from packaging.specifiers import SpecifierSet
from packaging.version import Version

from pants.backend.python.goals import lockfile
from pants.backend.python.goals.lockfile import GeneratePythonLockfile
from pants.backend.python.subsystems.setup import PythonSetup
from pants.backend.python.subsystems.uv import DownloadedUv
from pants.backend.python.target_types import EntryPoint, PexCompletePlatformsField
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,
    PexProcess,
    PexRequest,
    PexRequirementsInfo,
    PexResolveInfo,
    VenvPex,
    VenvPexProcess,
    _build_pex_description,
    _build_uv_venv,
    _BuildPexPythonSetup,
    _BuildPexRequirementsSetup,
    _determine_pex_python_and_platforms,
    _setup_pex_requirements,
    _UvVenvRequest,
)
from pants.backend.python.util_rules.pex import rules as pex_rules
from pants.backend.python.util_rules.pex_cli import PexCliProcess
from pants.backend.python.util_rules.pex_environment import PythonExecutable
from pants.backend.python.util_rules.pex_requirements import (
    EntireLockfile,
    LoadedLockfile,
    Lockfile,
    PexRequirements,
    Resolve,
    ResolvePexConfig,
)
from pants.backend.python.util_rules.pex_test_utils import (
    create_pex_and_get_all_data,
    create_pex_and_get_pex_info,
    parse_requirements,
)
from pants.core.goals.generate_lockfiles import GenerateLockfileResult
from pants.core.register import wrap_as_resources
from pants.core.target_types import FileTarget, ResourceTarget
from pants.core.util_rules.lockfile_metadata import InvalidLockfileError
from pants.engine.fs import (
    EMPTY_DIGEST,
    CreateDigest,
    Digest,
    DigestContents,
    Directory,
    FileContent,
)
from pants.engine.internals.native_engine import Address
from pants.engine.process import Process, ProcessCacheScope, ProcessResult
from pants.option.global_options import GlobalOptions
from pants.testutil.option_util import create_subsystem
from pants.testutil.rule_runner import (
    PYTHON_BOOTSTRAP_ENV,
    QueryRule,
    RuleRunner,
    engine_error,
    run_rule_with_mocks,
)
from pants.util.contextutil import temporary_dir
from pants.util.dirutil import safe_rmtree
from pants.util.ordered_set import FrozenOrderedSet
from pants.util.pip_requirement import PipRequirement


@pytest.fixture
def rule_runner() -> RuleRunner:
    return RuleRunner(
        rules=[
            *pex_test_utils.rules(),
            *pex_rules(),
            *wrap_as_resources.rules,
            QueryRule(GlobalOptions, []),
            QueryRule(ProcessResult, (Process,)),
            QueryRule(PexResolveInfo, (Pex,)),
            QueryRule(PexResolveInfo, (VenvPex,)),
            QueryRule(CompletePlatforms, (PexCompletePlatformsField,)),
        ],
        target_types=[
            ResourceTarget,
            FileTarget,
            *wrap_as_resources.target_types,
        ],
    )


@pytest.mark.parametrize("pex_type", [Pex, VenvPex])
@pytest.mark.parametrize("internal_only", [True, False])
def test_pex_execution(
    rule_runner: RuleRunner, pex_type: type[Pex | VenvPex], internal_only: bool
) -> None:
    sources = rule_runner.request(
        Digest,
        [
            CreateDigest(
                (
                    FileContent("main.py", b'print("from main")'),
                    FileContent("subdir/sub.py", b'print("from sub")'),
                )
            ),
        ],
    )
    pex_data = create_pex_and_get_all_data(
        rule_runner,
        pex_type=pex_type,
        internal_only=internal_only,
        main=EntryPoint("main"),
        sources=sources,
    )

    assert "main.py" in pex_data.files
    assert "subdir/sub.py" in pex_data.files

    # This should run the Pex using the same interpreter used to create it. We must set the `PATH`
    # so that the shebang works.
    pex_exe = (
        f"./{pex_data.sandbox_path}"
        if pex_data.is_zipapp
        else os.path.join(pex_data.sandbox_path, "__main__.py")
    )
    process = Process(
        argv=(pex_exe,),
        env={"PATH": os.getenv("PATH", "")},
        input_digest=pex_data.pex.digest,
        description="Run the pex and make sure it works",
    )
    result = rule_runner.request(ProcessResult, [process])
    assert result.stdout == b"from main\n"


@pytest.mark.parametrize("pex_type", [Pex, VenvPex])
def test_pex_environment(rule_runner: RuleRunner, pex_type: type[Pex | VenvPex]) -> None:
    sources = rule_runner.request(
        Digest,
        [
            CreateDigest(
                (
                    FileContent(
                        path="main.py",
                        content=textwrap.dedent(
                            """
                            from os import environ
                            print(f"LANG={environ.get('LANG')}")
                            print(f"ftp_proxy={environ.get('ftp_proxy')}")
                            """
                        ).encode(),
                    ),
                )
            ),
        ],
    )
    pex_data = create_pex_and_get_all_data(
        rule_runner,
        pex_type=pex_type,
        main=EntryPoint("main"),
        sources=sources,
        additional_pants_args=(
            "--subprocess-environment-env-vars=LANG",  # Value should come from environment.
            "--subprocess-environment-env-vars=ftp_proxy=dummyproxy",
        ),
        interpreter_constraints=InterpreterConstraints(["CPython>=3.8"]),
        env={"LANG": "es_PY.UTF-8"},
    )

    pex_process_type = PexProcess if isinstance(pex_data.pex, Pex) else VenvPexProcess
    process = rule_runner.request(
        Process,
        [
            pex_process_type(
                pex_data.pex,  # type: ignore[arg-type]
                description="Run the pex and check its reported environment",
            ),
        ],
    )

    result = rule_runner.request(ProcessResult, [process])
    assert b"LANG=es_PY.UTF-8" in result.stdout
    assert b"ftp_proxy=dummyproxy" in result.stdout


@pytest.mark.parametrize("pex_type", [Pex, VenvPex])
def test_pex_working_directory(rule_runner: RuleRunner, pex_type: type[Pex | VenvPex]) -> None:
    named_caches_dir = rule_runner.request(GlobalOptions, []).named_caches_dir
    sources = rule_runner.request(
        Digest,
        [
            CreateDigest(
                (
                    FileContent(
                        path="main.py",
                        content=textwrap.dedent(
                            """
                            import os
                            cwd = os.getcwd()
                            print(f"CWD: {cwd}")
                            for path, dirs, _ in os.walk(cwd):
                                for name in dirs:
                                    print(f"DIR: {os.path.relpath(os.path.join(path, name), cwd)}")
                            """
                        ).encode(),
                    ),
                )
            ),
        ],
    )

    pex_data = create_pex_and_get_all_data(
        rule_runner,
        pex_type=pex_type,
        main=EntryPoint("main"),
        sources=sources,
        interpreter_constraints=InterpreterConstraints(["CPython>=3.8"]),
    )

    pex_process_type = PexProcess if isinstance(pex_data.pex, Pex) else VenvPexProcess

    dirpath = "foo/bar/baz"
    runtime_files = rule_runner.request(Digest, [CreateDigest([Directory(path=dirpath)])])

    dirpath_parts = os.path.split(dirpath)
    for i in range(0, len(dirpath_parts)):
        working_dir = os.path.join(*dirpath_parts[:i]) if i > 0 else None
        expected_subdir = os.path.join(*dirpath_parts[i:]) if i < len(dirpath_parts) else None
        process = rule_runner.request(
            Process,
            [
                pex_process_type(
                    pex_data.pex,  # type: ignore[arg-type]
                    description="Run the pex and check its cwd",
                    working_directory=working_dir,
                    input_digest=runtime_files,
                    # We skip the process cache for this PEX to ensure that it re-runs.
                    cache_scope=ProcessCacheScope.PER_SESSION,
                )
            ],
        )

        # For VenvPexes, run the PEX twice while clearing the venv dir in between. This emulates
        # situations where a PEX creation hits the process cache, while venv seeding misses the PEX
        # cache.
        if isinstance(pex_data.pex, VenvPex):
            # Request once to ensure that the directory is seeded, and then start a new session so
            # that the second run happens as well.
            _ = rule_runner.request(ProcessResult, [process])
            rule_runner.new_session("re-run-for-venv-pex")
            rule_runner.set_options(
                ["--backend-packages=pants.backend.python"],
                env_inherit={"PATH", "PYENV_ROOT", "HOME"},
            )

            # Clear the cache.
            venv_dir = os.path.join(named_caches_dir, "pex_root", pex_data.pex.venv_rel_dir)
            assert os.path.isdir(venv_dir)
            safe_rmtree(venv_dir)

        result = rule_runner.request(ProcessResult, [process])
        output_str = result.stdout.decode()
        mo = re.search(r"CWD: (.*)\n", output_str)
        assert mo is not None
        reported_cwd = mo.group(1)
        if working_dir:
            assert reported_cwd.endswith(working_dir)
        if expected_subdir:
            assert f"DIR: {expected_subdir}" in output_str


def test_resolves_dependencies(rule_runner: RuleRunner) -> None:
    req_strings = ["six==1.12.0", "jsonschema==2.6.0", "requests==2.23.0"]
    requirements = PexRequirements(req_strings)
    pex_info = create_pex_and_get_pex_info(rule_runner, requirements=requirements)
    # NB: We do not check for transitive dependencies, which PEX-INFO will include. We only check
    # that at least the dependencies we requested are included.
    assert set(parse_requirements(req_strings)).issubset(
        set(parse_requirements(pex_info["requirements"]))
    )


def test_requirement_constraints(rule_runner: RuleRunner) -> None:
    direct_deps = ["requests>=1.0.0,<=2.23.0"]

    def assert_direct_requirements(pex_info):
        assert {PipRequirement.parse(r) for r in pex_info["requirements"]} == {
            PipRequirement.parse(d) for d in direct_deps
        }

    # Unconstrained, we should always pick the top of the range (requests 2.23.0) since the top of
    # the range is a transitive closure over universal wheels.
    direct_pex_info = create_pex_and_get_pex_info(
        rule_runner, requirements=PexRequirements(direct_deps)
    )
    assert_direct_requirements(direct_pex_info)
    assert "requests-2.23.0-py2.py3-none-any.whl" in set(direct_pex_info["distributions"].keys())

    constraints = [
        "requests==2.16.0",
        "certifi==2019.6.16",
        "chardet==3.0.2",
        "idna==2.5",
        "urllib3==1.21.1",
    ]
    rule_runner.write_files({"constraints.txt": "\n".join(constraints)})
    constrained_pex_info = create_pex_and_get_pex_info(
        rule_runner,
        requirements=PexRequirements(direct_deps, constraints_strings=constraints),
        additional_pants_args=("--python-requirement-constraints=constraints.txt",),
    )
    assert_direct_requirements(constrained_pex_info)
    assert {
        "certifi-2019.6.16-py2.py3-none-any.whl",
        "chardet-3.0.2-py2.py3-none-any.whl",
        "idna-2.5-py2.py3-none-any.whl",
        "requests-2.16.0-py2.py3-none-any.whl",
        "urllib3-1.21.1-py2.py3-none-any.whl",
    } == set(constrained_pex_info["distributions"].keys())


def test_lockfiles(rule_runner: RuleRunner) -> None:
    rule_runner.set_options(["--python-invalid-lockfile-behavior=ignore"])
    rule_runner.write_files(
        {
            "pex_lock.json": textwrap.dedent(
                """\
                // Some Pants header
                // blah blah
                {
                  "allow_builds": true,
                  "allow_prereleases": false,
                  "allow_wheels": true,
                  "build_isolation": true,
                  "constraints": [],
                  "locked_resolves": [
                    {
                      "locked_requirements": [
                        {
                          "artifacts": [
                            {
                              "algorithm": "sha256",
                              "hash": "00d2dde5a675579325902536738dd27e4fac1fd68f773fe36c21044eb559e187",
                              "url": "https://files.pythonhosted.org/packages/53/18/a56e2fe47b259bb52201093a3a9d4a32014f9d85071ad07e9d60600890ca/ansicolors-1.1.8-py2.py3-none-any.whl"
                            },
                            {
                              "algorithm": "sha256",
                              "hash": "99f94f5e3348a0bcd43c82e5fc4414013ccc19d70bd939ad71e0133ce9c372e0",
                              "url": "https://files.pythonhosted.org/packages/76/31/7faed52088732704523c259e24c26ce6f2f33fbeff2ff59274560c27628e/ansicolors-1.1.8.zip"
                            }
                          ],
                          "project_name": "ansicolors",
                          "requires_dists": [],
                          "requires_python": null,
                          "version": "1.1.8"
                        }
                      ],
                      "platform_tag": [
                        "cp39",
                        "cp39",
                        "macosx_11_0_arm64"
                      ]
                    }
                  ],
                  "pex_version": "2.1.70",
                  "prefer_older_binary": false,
                  "requirements": [
                    "ansicolors"
                  ],
                  "requires_python": [],
                  "resolver_version": "pip-2020-resolver",
                  "style": "universal",
                  "transitive": true,
                  "use_pep517": null
                }
                """
            ),
            "reqs_lock.txt": textwrap.dedent(
                """\
                ansicolors==1.1.8 \
                    --hash=sha256:00d2dde5a675579325902536738dd27e4fac1fd68f773fe36c21044eb559e187 \
                    --hash=sha256:99f94f5e3348a0bcd43c82e5fc4414013ccc19d70bd939ad71e0133ce9c372e0
                """
            ),
        }
    )

    def create_lock(path: str) -> None:
        lock = Lockfile(
            path,
            url_description_of_origin="foo",
            resolve_name="a",
        )
        create_pex_and_get_pex_info(
            rule_runner,
            requirements=EntireLockfile(lock, ("ansicolors",)),
            additional_pants_args=("--python-invalid-lockfile-behavior=ignore",),
        )

    create_lock("pex_lock.json")
    create_lock("reqs_lock.txt")


def test_entry_point(rule_runner: RuleRunner) -> None:
    entry_point = "pydoc"
    pex_info = create_pex_and_get_pex_info(rule_runner, main=EntryPoint(entry_point))
    assert pex_info["entry_point"] == entry_point


def test_interpreter_constraints(rule_runner: RuleRunner) -> None:
    constraints = InterpreterConstraints(["CPython>=2.7,<3", "CPython>=3.8,<3.12"])
    pex_info = create_pex_and_get_pex_info(
        rule_runner, interpreter_constraints=constraints, internal_only=False
    )
    assert set(pex_info["interpreter_constraints"]) == {str(c) for c in constraints}


def test_additional_args(rule_runner: RuleRunner) -> None:
    pex_info = create_pex_and_get_pex_info(rule_runner, additional_pex_args=("--no-strip-pex-env",))
    assert pex_info["strip_pex_env"] is False


def test_platforms(rule_runner: RuleRunner) -> None:
    # We use Python 2.7, rather than Python 3, to ensure that the specified platform is
    # actually used.
    platforms = PexPlatforms(["linux-x86_64-cp-27-cp27mu"])
    constraints = InterpreterConstraints(["CPython>=2.7,<3", "CPython>=3.8"])
    pex_data = create_pex_and_get_all_data(
        rule_runner,
        requirements=PexRequirements(["cryptography==2.9"]),
        platforms=platforms,
        interpreter_constraints=constraints,
        internal_only=False,  # Internal only PEXes do not support (foreign) platforms.
    )
    assert any(
        "cryptography-2.9-cp27-cp27mu-manylinux2010_x86_64.whl" in fp for fp in pex_data.files
    )
    assert not any("cryptography-2.9-cp27-cp27m-" in fp for fp in pex_data.files)
    assert not any("cryptography-2.9-cp35-abi3" in fp for fp in pex_data.files)

    # NB: Platforms override interpreter constraints.
    assert pex_data.info["interpreter_constraints"] == []


@pytest.mark.parametrize("use_pep440_rather_than_find_links", [True, False])
def test_local_requirements_and_path_mappings(
    use_pep440_rather_than_find_links: bool, tmp_path
) -> None:
    rule_runner = RuleRunner(
        rules=[
            *pex_test_utils.rules(),
            *pex_rules(),
            *lockfile.rules(),
            QueryRule(GenerateLockfileResult, [GeneratePythonLockfile]),
            QueryRule(PexResolveInfo, (Pex,)),
        ],
        bootstrap_args=[f"--named-caches-dir={tmp_path}"],
    )

    wheel_content = requests.get(
        "https://files.pythonhosted.org/packages/53/18/a56e2fe47b259bb52201093a3a9d4a32014f9d85071ad07e9d60600890ca/ansicolors-1.1.8-py2.py3-none-any.whl"
    ).content

    with temporary_dir() as wheel_base_dir:
        dir1_path = Path(wheel_base_dir, "dir1")
        dir2_path = Path(wheel_base_dir, "dir2")
        dir1_path.mkdir()
        dir2_path.mkdir()

        wheel_path = dir1_path / "ansicolors-1.1.8-py2.py3-none-any.whl"
        wheel_req_str = (
            f"ansicolors @ file://{wheel_path}"
            if use_pep440_rather_than_find_links
            else "ansicolors"
        )
        wheel_path.write_bytes(wheel_content)

        def options(path_mappings_dir: Path) -> tuple[str, ...]:
            return (
                "--python-repos-indexes=[]",
                (
                    "--python-repos-find-links=[]"
                    if use_pep440_rather_than_find_links
                    else f"--python-repos-find-links={path_mappings_dir}"
                ),
                f"--python-repos-path-mappings=WHEEL_DIR|{path_mappings_dir}",
                f"--named-caches-dir={tmp_path}",
                # Use the vendored pip, so we don't have to set up a wheel for it in dir1_path.
                "--python-pip-version=20.3.4-patched",
            )

        rule_runner.set_options(options(dir1_path), env_inherit=PYTHON_BOOTSTRAP_ENV)
        lock_result = rule_runner.request(
            GenerateLockfileResult,
            [
                GeneratePythonLockfile(
                    requirements=FrozenOrderedSet([wheel_req_str]),
                    find_links=FrozenOrderedSet([]),
                    interpreter_constraints=InterpreterConstraints([">=3.8,<3.15"]),
                    resolve_name="test",
                    lockfile_dest="test.lock",
                    diff=False,
                    lock_style="universal",
                    complete_platforms=(),
                )
            ],
        )
        lock_digest_contents = rule_runner.request(DigestContents, [lock_result.digest])
        assert len(lock_digest_contents) == 1
        lock_file_content = lock_digest_contents[0]
        assert b"${WHEEL_DIR}/ansicolors-1.1.8-py2.py3-none-any.whl" in lock_file_content.content
        assert b"files.pythonhosted.org" not in lock_file_content.content

        rule_runner.write_files({"test.lock": lock_file_content.content})
        lockfile_obj = EntireLockfile(
            Lockfile(url="test.lock", url_description_of_origin="test", resolve_name="test"),
            (wheel_req_str,),
        )

        # Wipe cache to ensure `--path-mappings` works.
        shutil.rmtree(tmp_path)
        shutil.rmtree(dir1_path)
        (dir2_path / "ansicolors-1.1.8-py2.py3-none-any.whl").write_bytes(wheel_content)
        pex_info = create_pex_and_get_all_data(
            rule_runner, requirements=lockfile_obj, additional_pants_args=options(dir2_path)
        ).info
        assert "ansicolors-1.1.8-py2.py3-none-any.whl" in pex_info["distributions"]

        # Confirm that pointing to a bad path fails.
        shutil.rmtree(tmp_path)
        shutil.rmtree(dir2_path)
        with engine_error():
            create_pex_and_get_all_data(
                rule_runner,
                requirements=lockfile_obj,
                additional_pants_args=options(Path(wheel_base_dir, "dir3")),
            )


@pytest.mark.parametrize("pex_type", [Pex, VenvPex])
@pytest.mark.parametrize("internal_only", [True, False])
def test_additional_inputs(
    rule_runner: RuleRunner, pex_type: type[Pex | VenvPex], internal_only: bool
) -> None:
    # We use Pex's --sources-directory option to add an extra source file to the PEX.
    # This verifies that the file was indeed provided as additional input to the pex call.
    extra_src_dir = "extra_src"
    data_file = os.path.join("data", "file")
    data = "42"
    additional_inputs = rule_runner.request(
        Digest,
        [
            CreateDigest(
                [FileContent(path=os.path.join(extra_src_dir, data_file), content=data.encode())]
            )
        ],
    )
    additional_pex_args = ("--sources-directory", extra_src_dir)
    pex_data = create_pex_and_get_all_data(
        rule_runner,
        pex_type=pex_type,
        internal_only=internal_only,
        additional_inputs=additional_inputs,
        additional_pex_args=additional_pex_args,
    )
    if pex_data.is_zipapp:
        with zipfile.ZipFile(pex_data.local_path, "r") as zipfp:
            with zipfp.open(data_file, "r") as datafp:
                data_file_content = datafp.read()
    else:
        with open(pex_data.local_path / data_file, "rb") as datafp:
            data_file_content = datafp.read()
    assert data == data_file_content.decode()


@pytest.mark.parametrize("pex_type", [Pex, VenvPex])
def test_venv_pex_resolve_info(rule_runner: RuleRunner, pex_type: type[Pex | VenvPex]) -> None:
    constraints = [
        "requests==2.23.0",
        "certifi==2020.12.5",
        "chardet==3.0.4",
        "idna==2.10",
        "urllib3==1.25.11",
    ]
    rule_runner.write_files({"constraints.txt": "\n".join(constraints)})
    pex = create_pex_and_get_all_data(
        rule_runner,
        pex_type=pex_type,
        requirements=PexRequirements(["requests==2.23.0"], constraints_strings=constraints),
        additional_pants_args=("--python-requirement-constraints=constraints.txt",),
    ).pex
    dists = rule_runner.request(PexResolveInfo, [pex])
    assert dists[0] == PexDistributionInfo("certifi", Version("2020.12.5"), None, ())
    assert dists[1] == PexDistributionInfo("chardet", Version("3.0.4"), None, ())
    assert dists[2] == PexDistributionInfo(
        "idna", Version("2.10"), SpecifierSet("!=3.0.*,!=3.1.*,!=3.2.*,!=3.3.*,>=2.7"), ()
    )
    assert dists[3].project_name == "requests"
    assert dists[3].version == Version("2.23.0")
    # requires_dists is parsed from metadata written by the pex tool, and is always
    #   a set of valid requirements.
    assert Requirement('PySocks!=1.5.7,>=1.5.6; extra == "socks"') in dists[3].requires_dists
    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.8"])

    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_calls={
                "pants.backend.python.util_rules.pex.find_interpreter": lambda _: discovered_python
            },
        )
        assert result == expected

    assert_setup(
        expected=_BuildPexPythonSetup(discovered_python, ["--python-path", discovered_python.path])
    )
    assert_setup(
        interpreter_constraints=ics,
        expected=_BuildPexPythonSetup(
            discovered_python,
            ["--interpreter-constraint", "CPython==3.8", "--python-path", discovered_python.path],
        ),
    )
    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, url_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,
            as_constraints_strings=None,
            original_lockfile=lockfile_obj,
        )

    def assert_setup(
        requirements: PexRequirements | EntireLockfile,
        expected: _BuildPexRequirementsSetup,
        *,
        is_pex_lock: bool = True,
        include_find_links: bool = False,
    ) -> 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(PythonSetup)],
            mock_calls={
                "pants.backend.python.util_rules.pex_requirements.determine_resolve_pex_config": lambda _: ResolvePexConfig(
                    indexes=("custom-index",),
                    find_links=("custom-find-links",),
                    manylinux=None,
                    constraints_file=None,
                    only_binary=FrozenOrderedSet(),
                    no_binary=FrozenOrderedSet(),
                    path_mappings=(),
                    excludes=FrozenOrderedSet(),
                    overrides=FrozenOrderedSet(),
                    sources=FrozenOrderedSet(),
                    lock_style="universal",
                    complete_platforms=(),
                    uploaded_prior_to=None,
                ),
                "pants.backend.python.util_rules.pex.get_req_strings": lambda _: PexRequirementsInfo(
                    (
                        tuple(str(x) for x in requirements.req_strings_or_addrs)
                        if isinstance(requirements, PexRequirements)
                        else tuple()
                    ),
                    ("imma/link",) if include_find_links else tuple(),
                ),
                "pants.engine.intrinsics.create_digest": lambda _: constraints_digest,
                "pants.backend.python.util_rules.pex_requirements.load_lockfile": lambda _: create_loaded_lockfile(
                    is_pex_lock
                ),
                "pants.backend.python.util_rules.pex_requirements.get_lockfile_for_resolve": lambda _: lockfile_obj,
            },
        )
        assert result == expected

    pex_args = [
        "--no-pypi",
        "--index=custom-index",
        "--find-links=custom-find-links",
        "--no-manylinux",
    ]
    pip_args = [*pex_args, "--resolver-version", "pip-2020-resolver"]

    # 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(
            [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=Resolve("resolve", False)),
        _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(rule_runner: RuleRunner) -> None:
    def assert_description(
        requirements: PexRequirements | EntireLockfile,
        *,
        description: str | None = None,
        expected: str,
    ) -> None:
        request = PexRequest(
            output_filename="new.pex",
            internal_only=True,
            requirements=requirements,
            description=description,
        )
        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)

    assert_description(PexRequirements(), description="Custom!", expected="Custom!")
    assert_description(
        PexRequirements(from_superset=repo_pex), description="Custom!", expected="Custom!"
    )

    assert_description(PexRequirements(), expected="Building new.pex")
    assert_description(PexRequirements(from_superset=repo_pex), expected="Building new.pex")

    assert_description(
        PexRequirements(["req"]), expected="Building new.pex with 1 requirement: req"
    )
    assert_description(
        PexRequirements(["req"], from_superset=repo_pex),
        expected="Extracting 1 requirement to build new.pex from repo.pex: req",
    )

    assert_description(
        PexRequirements(["req1", "req2"]),
        expected="Building new.pex with 2 requirements: req1, req2",
    )
    assert_description(
        PexRequirements(["req1", "req2"], from_superset=repo_pex),
        expected="Extracting 2 requirements to build new.pex from repo.pex: req1, req2",
    )

    assert_description(
        EntireLockfile(
            Lockfile(
                url="lock.txt",
                url_description_of_origin="test",
                resolve_name="a",
            )
        ),
        expected="Building new.pex from lock.txt",
    )

    assert_description(
        EntireLockfile(
            Lockfile(
                url="lock.txt",
                url_description_of_origin="foo",
                resolve_name="a",
            )
        ),
        expected="Building new.pex from lock.txt",
    )


def test_lockfile_validation(rule_runner: RuleRunner) -> None:
    """Check that we properly load and validate lockfile metadata for both types of locks.

    Note that we don't exhaustively test every source of lockfile failure nor the different options
    for `--invalid-lockfile-behavior`, as those are already tested in pex_requirements_test.py.
    """

    # We create a lockfile that claims it works with no requirements. It should fail when we try
    # to build a PEX with a requirement.
    lock_content = PythonLockfileMetadata.new(
        valid_for_interpreter_constraints=InterpreterConstraints(),
        requirements=set(),
        requirement_constraints=set(),
        only_binary=set(),
        no_binary=set(),
        manylinux=None,
        excludes=set(),
        overrides=set(),
        sources=set(),
        lock_style="universal",
        complete_platforms=(),
        uploaded_prior_to=None,
    ).add_header_to_lockfile(b"", regenerate_command="regen", delimeter="#")
    rule_runner.write_files({"lock.txt": lock_content.decode()})

    _lockfile = Lockfile(
        "lock.txt",
        url_description_of_origin="a test",
        resolve_name="a",
    )
    with engine_error(InvalidLockfileError):
        create_pex_and_get_all_data(
            rule_runner, requirements=EntireLockfile(_lockfile, ("ansicolors",))
        )


@pytest.mark.parametrize("target_type", ["file", "resource"])
def test_digest_complete_platforms(rule_runner: RuleRunner, target_type: str) -> None:
    # Read the complete_platforms content using pkgutil
    complete_platforms_content = pkgutil.get_data(__name__, "complete_platform_pex_test.json")
    assert complete_platforms_content is not None

    # Create a target with the complete platforms file
    rule_runner.write_files(
        {
            "BUILD": f"{target_type}(name='complete_platforms', source='complete_platforms.json')",
            "complete_platforms.json": complete_platforms_content,
        }
    )

    # Get the CompletePlatforms object
    target = rule_runner.get_target(Address("", target_name="complete_platforms"))
    complete_platforms = rule_runner.request(
        CompletePlatforms,
        [PexCompletePlatformsField([":complete_platforms"], target.address)],
    )

    # Verify the result
    assert len(complete_platforms) == 1
    assert complete_platforms.digest != EMPTY_DIGEST


def test_digest_complete_platforms_codegen(rule_runner: RuleRunner) -> None:
    # Read the complete_platforms content using pkgutil
    complete_platforms_content = pkgutil.get_data(__name__, "complete_platform_pex_test.json")
    assert complete_platforms_content is not None

    # Create a target with the complete platforms file
    rule_runner.write_files(
        {
            "BUILD": """\
file(name='complete_platforms', source='complete_platforms.json')
experimental_wrap_as_resources(name="codegen", inputs=[':complete_platforms'], )
            """,
            "complete_platforms.json": complete_platforms_content,
        }
    )

    # Get the CompletePlatforms object
    target = rule_runner.get_target(Address("", target_name="codegen"))
    complete_platforms = rule_runner.request(
        CompletePlatforms,
        [PexCompletePlatformsField([":codegen"], target.address)],
    )

    # Verify the result
    assert len(complete_platforms) == 1
    assert complete_platforms.digest != EMPTY_DIGEST


def create_uv_venv_test_inputs(
    rule_runner: RuleRunner,
    *,
    description: str,
) -> tuple[_UvVenvRequest, Lockfile, LoadedLockfile]:
    resolve = Resolve("python-default", False)
    lockfile = Lockfile(
        "3rdparty/python/default.lock", url_description_of_origin="test", resolve_name=resolve.name
    )
    loaded_lockfile = LoadedLockfile(
        lockfile_digest=rule_runner.make_snapshot_of_empty_files([lockfile.url]).digest,
        lockfile_path=lockfile.url,
        metadata=None,
        requirement_estimate=1,
        is_pex_native=True,
        as_constraints_strings=None,
        original_lockfile=lockfile,
    )
    uv_request = _UvVenvRequest(
        req_strings=("ansicolors==1.1.8",),
        requirements=PexRequirements(("ansicolors==1.1.8",), from_superset=resolve),
        python_path="/usr/bin/python3",
        description=description,
    )
    return uv_request, lockfile, loaded_lockfile


def run_uv_venv_with_mocks(
    uv_request: _UvVenvRequest,
    lockfile: Lockfile,
    loaded_lockfile: LoadedLockfile,
    *,
    mock_fallible_to_exec_result_or_raise,
    mock_create_digest=None,
):
    pex_env = SimpleNamespace(append_only_caches={})
    pex_env.in_sandbox = lambda *, working_directory: pex_env
    pex_env.environment_dict = lambda *, python_configured: {}

    return run_rule_with_mocks(
        _build_uv_venv,
        rule_args=[uv_request, pex_env],
        mock_calls={
            "pants.backend.python.subsystems.uv.download_uv_binary": lambda: DownloadedUv(
                digest=EMPTY_DIGEST,
                exe="uv",
                args_for_uv_pip_install=(),
            ),
            "pants.backend.python.util_rules.pex_requirements.get_lockfile_for_resolve": lambda _: lockfile,
            "pants.backend.python.util_rules.pex_requirements.load_lockfile": lambda _: loaded_lockfile,
            "pants.engine.process.fallible_to_exec_result_or_raise": mock_fallible_to_exec_result_or_raise,
            "pants.engine.intrinsics.create_digest": mock_create_digest or (lambda _: EMPTY_DIGEST),
            "pants.engine.intrinsics.merge_digests": lambda _: EMPTY_DIGEST,
        },
    )


def test_build_uv_venv_uses_exported_lockfile_with_no_deps(rule_runner: RuleRunner) -> None:
    uv_request, lockfile, loaded_lockfile = create_uv_venv_test_inputs(
        rule_runner, description="test uv export path"
    )
    exported_digest = rule_runner.make_snapshot(
        {"__uv_requirements.txt": "ansicolors==1.1.8\n"}
    ).digest
    install_argv: tuple[str, ...] | None = None

    def mock_fallible_to_exec_result_or_raise(*args, **kwargs):
        nonlocal install_argv
        req = args[0]
        if isinstance(req, PexCliProcess):
            assert req.subcommand == ("lock", "export-subset")
            assert "--format" in req.extra_args
            export_format = req.extra_args[req.extra_args.index("--format") + 1]
            assert export_format in {"pip-no-hashes", "pep-751"}
            # Verify requirement strings are passed as positional args.
            assert "ansicolors==1.1.8" in req.extra_args
            # Verify --lock is used instead of positional lockfile path.
            assert "--lock" in req.extra_args
            return SimpleNamespace(output_digest=exported_digest)
        assert isinstance(req, Process)
        if req.argv[1:3] == ("pip", "install"):
            install_argv = req.argv
        return SimpleNamespace(output_digest=EMPTY_DIGEST)

    def mock_create_digest(request: CreateDigest) -> Digest:
        for entry in request:
            assert not (isinstance(entry, FileContent) and entry.path == "__uv_requirements.txt"), (
                "exported lockfile path should not synthesize requirements content"
            )
        return EMPTY_DIGEST

    result = run_uv_venv_with_mocks(
        uv_request,
        lockfile,
        loaded_lockfile,
        mock_fallible_to_exec_result_or_raise=mock_fallible_to_exec_result_or_raise,
        mock_create_digest=mock_create_digest,
    )

    assert result.venv_digest == EMPTY_DIGEST
    assert install_argv is not None
    assert "--no-deps" in install_argv


def test_build_uv_venv_falls_back_when_lock_export_has_no_digest(rule_runner: RuleRunner) -> None:
    uv_request, lockfile, loaded_lockfile = create_uv_venv_test_inputs(
        rule_runner, description="test uv fallback path"
    )
    export_attempts = 0
    install_argv: tuple[str, ...] | None = None
    synthesized_reqs: bytes | None = None

    def mock_fallible_to_exec_result_or_raise(*args, **kwargs):
        nonlocal export_attempts, install_argv
        req = args[0]
        if isinstance(req, PexCliProcess):
            export_attempts += 1
            return SimpleNamespace(output_digest=None)
        assert isinstance(req, Process)
        if req.argv[1:3] == ("pip", "install"):
            install_argv = req.argv
        return SimpleNamespace(output_digest=EMPTY_DIGEST)

    def mock_create_digest(request: CreateDigest) -> Digest:
        nonlocal synthesized_reqs
        for entry in request:
            if isinstance(entry, FileContent) and entry.path == "__uv_requirements.txt":
                synthesized_reqs = entry.content
        return EMPTY_DIGEST

    result = run_uv_venv_with_mocks(
        uv_request,
        lockfile,
        loaded_lockfile,
        mock_fallible_to_exec_result_or_raise=mock_fallible_to_exec_result_or_raise,
        mock_create_digest=mock_create_digest,
    )

    assert result.venv_digest == EMPTY_DIGEST
    assert export_attempts == 1
    assert install_argv is not None
    assert "--no-deps" not in install_argv
    assert synthesized_reqs == b"ansicolors==1.1.8\n"


def test_build_uv_venv_propagates_unexpected_export_errors(rule_runner: RuleRunner) -> None:
    uv_request, lockfile, loaded_lockfile = create_uv_venv_test_inputs(
        rule_runner, description="test unexpected error path"
    )

    def mock_fallible_to_exec_result_or_raise(*args, **kwargs):
        req = args[0]
        if isinstance(req, PexCliProcess):
            raise ValueError("unexpected failure type")
        return SimpleNamespace(output_digest=EMPTY_DIGEST)

    with pytest.raises(ValueError, match="unexpected failure type"):
        run_uv_venv_with_mocks(
            uv_request,
            lockfile,
            loaded_lockfile,
            mock_fallible_to_exec_result_or_raise=mock_fallible_to_exec_result_or_raise,
        )


def test_uv_pex_builder_resolves_dependencies(rule_runner: RuleRunner) -> None:
    """When pex_builder=uv, PEX should be built via uv venv + --venv-repository."""
    req_strings = ["six==1.12.0", "jsonschema==2.6.0"]
    requirements = PexRequirements(req_strings)
    pex_info = create_pex_and_get_pex_info(
        rule_runner,
        requirements=requirements,
        additional_pants_args=("--python-pex-builder=uv",),
        internal_only=False,
    )
    assert set(parse_requirements(req_strings)).issubset(
        set(parse_requirements(pex_info["requirements"]))
    )


def test_uv_pex_builder_includes_transitive_dependencies(rule_runner: RuleRunner) -> None:
    """uv builder must install transitive dependencies, not just direct ones.

    `requests` depends on urllib3, certifi, charset-normalizer, idna - the PEX
    must be able to import these at runtime even though only `requests` is declared.
    We verify by actually executing the PEX and importing a transitive dep.
    """
    sources = rule_runner.request(
        Digest,
        [
            CreateDigest(
                (
                    FileContent(
                        "main.py",
                        # Import both the direct dep and a transitive dep (certifi).
                        b"import requests; import certifi; print(f'requests=={requests.__version__}'); print(f'certifi_where={certifi.where()}')",
                    ),
                )
            ),
        ],
    )
    pex_data = create_pex_and_get_all_data(
        rule_runner,
        pex_type=Pex,
        requirements=PexRequirements(["requests==2.31.0"]),
        main=EntryPoint("main"),
        sources=sources,
        additional_pants_args=("--python-pex-builder=uv",),
        internal_only=False,
    )
    pex_exe = (
        f"./{pex_data.sandbox_path}"
        if pex_data.is_zipapp
        else os.path.join(pex_data.sandbox_path, "__main__.py")
    )
    process = Process(
        argv=(pex_exe,),
        env={"PATH": os.getenv("PATH", "")},
        input_digest=pex_data.pex.digest,
        description="Run uv-built pex and verify transitive deps are importable",
    )
    result = rule_runner.request(ProcessResult, [process])
    assert b"requests==2.31.0" in result.stdout
    assert b"certifi_where=" in result.stdout


def test_uv_pex_builder_execution(rule_runner: RuleRunner) -> None:
    """PEX built via uv builder should actually execute and import installed packages."""
    sources = rule_runner.request(
        Digest,
        [
            CreateDigest(
                (
                    FileContent(
                        "main.py",
                        b"import six; print(f'six=={six.__version__}')",
                    ),
                )
            ),
        ],
    )
    pex_data = create_pex_and_get_all_data(
        rule_runner,
        pex_type=Pex,
        requirements=PexRequirements(["six==1.12.0"]),
        main=EntryPoint("main"),
        sources=sources,
        additional_pants_args=("--python-pex-builder=uv",),
        internal_only=False,
    )
    pex_exe = (
        f"./{pex_data.sandbox_path}"
        if pex_data.is_zipapp
        else os.path.join(pex_data.sandbox_path, "__main__.py")
    )
    process = Process(
        argv=(pex_exe,),
        env={"PATH": os.getenv("PATH", "")},
        input_digest=pex_data.pex.digest,
        description="Run uv-built pex and verify import works",
    )
    result = rule_runner.request(ProcessResult, [process])
    assert result.stdout == b"six==1.12.0\n"


def test_uv_pex_builder_skipped_for_internal_only(rule_runner: RuleRunner) -> None:
    """Internal-only PEXes should fall back to the default pip path even with pex_builder=uv."""
    req_strings = ["six==1.12.0"]
    requirements = PexRequirements(req_strings)
    pex_info = create_pex_and_get_pex_info(
        rule_runner,
        requirements=requirements,
        additional_pants_args=("--python-pex-builder=uv",),
        internal_only=True,
    )
    assert set(parse_requirements(req_strings)).issubset(
        set(parse_requirements(pex_info["requirements"]))
    )

pantsbuild / pants / 24149123771

Source File Press 'n' to go to next uncovered line, 'b' for previous

Source File
Press 'n' to go to next uncovered line, 'b' for previous