From b96419c149a200de8df0b37cd797643b15c769c3 Mon Sep 17 00:00:00 2001 From: Benjy Weinberger Date: Tue, 5 May 2026 21:41:59 -0700 Subject: [PATCH] Add features to uv support. --- .../pants/backend/python/subsystems/uv.py | 13 ++++ .../python/util_rules/pex_requirements.py | 74 ++++++++++++++----- .../util_rules/pex_requirements_test.py | 34 +++++++-- .../pants/backend/python/util_rules/uv.py | 7 ++ 4 files changed, 104 insertions(+), 24 deletions(-) diff --git a/src/python/pants/backend/python/subsystems/uv.py b/src/python/pants/backend/python/subsystems/uv.py index 2e5aea6e64f..8630720d888 100644 --- a/src/python/pants/backend/python/subsystems/uv.py +++ b/src/python/pants/backend/python/subsystems/uv.py @@ -11,6 +11,7 @@ TemplatedExternalTool, download_external_tool, ) +from pants.engine.env_vars import EXTRA_ENV_VARS_USAGE_HELP from pants.engine.fs import Digest from pants.engine.internals.native_engine import FrozenDict from pants.engine.platform import Platform @@ -65,6 +66,18 @@ class EnvironmentAware(Subsystem.EnvironmentAware): metavar="", ) + extra_env_vars = StrListOption( + help=softwrap( + f""" + Additional environment variables to pass to `uv` subprocesses. + Can be used to pass `UV_KEYRING_PROVIDER`, or private-index credentials such as + `UV_INDEX__USERNAME`/`UV_INDEX__PASSWORD`. + + {EXTRA_ENV_VARS_USAGE_HELP} + """ + ), + ) + @memoized_property def path(self) -> tuple[str, ...]: def iter_path_entries(): diff --git a/src/python/pants/backend/python/util_rules/pex_requirements.py b/src/python/pants/backend/python/util_rules/pex_requirements.py index 652747c211e..64923102a08 100644 --- a/src/python/pants/backend/python/util_rules/pex_requirements.py +++ b/src/python/pants/backend/python/util_rules/pex_requirements.py @@ -12,6 +12,8 @@ from typing import TYPE_CHECKING from urllib.parse import urlparse +from packaging.requirements import Requirement + from pants.backend.python.subsystems.repos import PythonRepos from pants.backend.python.subsystems.setup import InvalidLockfileBehavior, PythonSetup from pants.backend.python.target_types import PythonRequirementsField @@ -488,37 +490,75 @@ def uv_config(self, extra_find_links: Iterable[str] = ()) -> str: """ config_lines: list[str] = [] - if not self.indexes: - config_lines.append("no-index = true") - all_find_links = (*self.find_links, *extra_find_links) if all_find_links: - entries = "\n".join(f' "{fl}",' for fl in all_find_links) - config_lines.append(f"find-links = [\n{entries}\n]") + config_lines.append("find-links = [") + for fl in all_find_links: + config_lines.append(f' "{fl}",') + config_lines.append("]") + config_lines.append("") + + if self.sources: + config_lines.append("[sources]") + for source in self.sources: + index_name, _, scope = source.partition("=") + req = Requirement(scope) + # Markers may contain double-quotes, so we use single quotes in the TOML. + marker = f", marker = '{req.marker}'" if req.marker else "" + config_lines.append(f'{req.name} = {{ index = "{index_name}"{marker} }}') + config_lines.append("") if self.no_binary: if ":all:" in self.no_binary: config_lines.append("no-binary = true") elif ":none:" not in self.no_binary: - entries = "\n".join(f' "{pkg}",' for pkg in self.no_binary) - config_lines.append(f"no-binary-package = [\n{entries}\n]") + config_lines.append("no-binary-package = [") + for pkg in self.no_binary: + config_lines.append(f' "{pkg}",') + config_lines.append("]") + config_lines.append("") if self.only_binary: if ":all:" in self.only_binary: config_lines.append("no-build = true") elif ":none:" not in self.only_binary: - entries = "\n".join(f' "{pkg}",' for pkg in self.only_binary) - config_lines.append(f"no-build-package = [\n{entries}\n]") + config_lines.append("no-build-package = [") + for pkg in self.only_binary: + config_lines.append(f' "{pkg}",') + config_lines.append("]") + config_lines.append("") if self.uploaded_prior_to: config_lines.append(f'exclude-newer = "{self.uploaded_prior_to}"') - - for i, index_url in enumerate(self.indexes): - # The first index gets `default = true`, replacing uv's built-in PyPI default. - # Subsequent indexes are additional sources. - block = f'[[index]]\nurl = "{index_url}"\n' - block += "default = true\n" if i == 0 else f'name = "extra-{i}"\n' - config_lines.append(block) + config_lines.append("") + + indexes = [] + for index in self.indexes: + part1, _, part2 = index.partition("=") + (name, url) = (part1, part2) if part2 else ("", part1) + index_data = {"url": url} + if name: + index_data["name"] = name + indexes.append(index_data) + if indexes: + # To turn off uv's fallback to PyPI we must set some other index to be the default. + # In uv the default index has the lowest priority, regardless of its position in the + # list of indexes, so we set the last index to be that default, to match user intent. + indexes[-1]["default"] = "true" + for index_data in indexes: + name = index_data.get("name", "") + url = index_data.get("url", "") + default = index_data.get("default", False) + config_lines.append("[[index]]") + if name: + config_lines.append(f'name = "{name}"') + config_lines.append(f'url = "{url}"') + if default: + config_lines.append("default = true") + config_lines.append("") + else: + config_lines.append("no-index = true") + config_lines.append("") return "\n".join(config_lines) + "\n" if config_lines else "" @@ -533,8 +573,6 @@ def validate_for_uv(self, resolve_name: str) -> None: pex_specific.append("`[python].resolves_to_excludes`") if self.overrides: pex_specific.append("`[python].resolves_to_overrides`") - if self.sources: - pex_specific.append("`[python].resolves_to_sources`") if self.lock_style != "universal": pex_specific.append("`[python]._resolves_to_lock_style`") if self.path_mappings: diff --git a/src/python/pants/backend/python/util_rules/pex_requirements_test.py b/src/python/pants/backend/python/util_rules/pex_requirements_test.py index 4ab6fb5c46c..5871f8bd124 100644 --- a/src/python/pants/backend/python/util_rules/pex_requirements_test.py +++ b/src/python/pants/backend/python/util_rules/pex_requirements_test.py @@ -457,6 +457,7 @@ def _uv_config( indexes=None, find_links=None, extra_find_links=(), + sources=None, only_binary=None, no_binary=None, uploaded_prior_to=None, @@ -466,11 +467,11 @@ def _uv_config( find_links=find_links or [], manylinux=None, constraints_file=None, - no_binary=FrozenOrderedSet(no_binary) if no_binary else FrozenOrderedSet(), - only_binary=FrozenOrderedSet(only_binary) if only_binary else FrozenOrderedSet(), + no_binary=FrozenOrderedSet(no_binary or []), + only_binary=FrozenOrderedSet(only_binary or []), excludes=FrozenOrderedSet(), overrides=FrozenOrderedSet(), - sources=FrozenOrderedSet(), + sources=FrozenOrderedSet(sources or []), path_mappings=[], lock_style="universal", complete_platforms=(), @@ -490,14 +491,19 @@ def test_uv_config_indexes(): assert parsed["index"][0]["default"] is True parsed = _uv_config( - indexes=["https://primary.example.com/simple", "https://secondary.example.com/simple"] + indexes=[ + "https://primary.example.com/simple", + "fallback=https://secondary.example.com/simple", + ] ) indexes = parsed["index"] assert len(indexes) == 2 assert indexes[0]["url"] == "https://primary.example.com/simple" - assert indexes[0]["default"] is True + assert "name" not in indexes[0] + assert "default" not in indexes[0] assert indexes[1]["url"] == "https://secondary.example.com/simple" - assert indexes[1].get("name") == "extra-1" + assert indexes[1].get("name") == "fallback" + assert indexes[1]["default"] is True def test_uv_config_find_links(): @@ -514,6 +520,22 @@ def test_uv_config_find_links(): ] +def test_uv_config_sources(): + parsed = _uv_config(sources=[]) + assert "sources" not in parsed + + parsed = _uv_config(sources=["myindex=requests>=2.0"]) + assert parsed["sources"] == {"requests": {"index": "myindex"}} + + parsed = _uv_config(sources=['myindex=requests>=2.0; python_version > "3.8"']) + assert parsed["sources"]["requests"]["index"] == "myindex" + assert "python_version" in parsed["sources"]["requests"]["marker"] + + parsed = _uv_config(sources=["indexa=requests>=2.0", "indexb=boto3>=1.0"]) + assert parsed["sources"]["requests"] == {"index": "indexa"} + assert parsed["sources"]["boto3"] == {"index": "indexb"} + + def test_uv_config_no_binary(): parsed = _uv_config(no_binary=["foo", "bar"]) assert parsed["no-binary-package"] == ["foo", "bar"] diff --git a/src/python/pants/backend/python/util_rules/uv.py b/src/python/pants/backend/python/util_rules/uv.py index 8ee58836d72..8f24cafb7cf 100644 --- a/src/python/pants/backend/python/util_rules/uv.py +++ b/src/python/pants/backend/python/util_rules/uv.py @@ -29,8 +29,10 @@ ) from pants.base.build_root import BuildRoot from pants.core.util_rules import system_binaries +from pants.core.util_rules.env_vars import environment_vars_subset from pants.core.util_rules.subprocess_environment import SubprocessEnvironmentVars from pants.core.util_rules.system_binaries import BashBinary, RealpathBinary +from pants.engine.env_vars import EnvironmentVarsRequest from pants.engine.fs import ( CreateDigest, FileContent, @@ -94,11 +96,16 @@ async def get_uv_environment( path = os.pathsep.join(uv_env_aware.path) subprocess_env_dict = dict(subprocess_env_vars.vars) + extra_env = await environment_vars_subset( + EnvironmentVarsRequest(uv_env_aware.extra_env_vars), **implicitly() + ) + if "PATH" in subprocess_env_dict: path = os.pathsep.join([path, subprocess_env_dict.pop("PATH")]) return UvEnvironment( env=FrozenDict( { + **extra_env, "PATH": path, **subprocess_env_dict, **python_native_code.subprocess_env_vars,