Add features to uv support.

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="<binary-paths>",
         )
 
+        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_<NAME>_USERNAME`/`UV_INDEX_<NAME>_PASSWORD`.
+
+                {EXTRA_ENV_VARS_USAGE_HELP}
+                """
+            ),
+        )
+
         @memoized_property
         def path(self) -> tuple[str, ...]:
             def iter_path_entries():

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:

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"]

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,