feat(uv): support test-only Python packages

README.md

@@ -207,6 +207,13 @@ uv.override_package(
     exclude_glob = ["another_package/**/tests/**"],
 )
 
+# 3d. (Optional) Restrict a package and its dependents to test targets
+uv.override_package(
+    name = "pytest-postgresql",
+    lock = "//:uv.lock",
+    testonly = True,
+)
+
 use_repo(uv, "pypi")
 ```
 

e2e/cases/uv-patching-829/BUILD.bazel

@@ -1,5 +1,12 @@
 load("@aspect_rules_py//py:defs.bzl", "py_test")
 
+filegroup(
+    name = "testonly_data",
+    testonly = True,
+    srcs = ["//uv-patching-829/patches:post_install.patch"],
+    visibility = ["//visibility:public"],
+)
+
 py_test(
     name = "test",
     srcs = ["__test__.py"],

e2e/cases/uv-patching-829/setup.MODULE.bazel

@@ -7,6 +7,8 @@ uv.project(
 )
 uv.override_package(
     name = "cowsay",
+    testonly = True,
+    extra_data = ["//uv-patching-829:testonly_data"],
     lock = "//uv-patching-829:uv.lock",
     post_install_patch_strip = 1,
     post_install_patches = ["//uv-patching-829/patches:post_install.patch"],

uv/private/extension/defs.bzl

@@ -98,6 +98,7 @@ def shared_install_key(install_cfg):
         # ordering is significant and the pairs stay as an ordered list.
         install_cfg.whls.items(),
         install_cfg.exclude_glob,
+        install_cfg.testonly,
     ])
 
 def dedupe_shared_installs(install_cfgs):
@@ -246,6 +247,7 @@ def _parse_projects(module_ctx, hub_specs):
 
             has_target = override.target != None
             has_modifications = (
+                override.testonly or
                 override.console_scripts != _CONSOLE_SCRIPTS_UNSET or
                 override.pre_build_patches or
                 override.post_install_patches or
@@ -260,7 +262,7 @@ def _parse_projects(module_ctx, hub_specs):
             if has_target and has_modifications:
                 fail("uv.override_package() for '{}': `target` is mutually exclusive with modification attributes. Use `target` for full replacement OR build, patch, and data attributes for modifications, not both.".format(override.name))
             if not has_target and not has_modifications:
-                fail("uv.override_package() for '{}': must specify either `target` for full replacement or at least one modification attribute (console_scripts, pre_build_patches, post_install_patches, exclude_glob, extra_deps, extra_data, toolchains, env, monitor_memory, resource_set).".format(override.name))
+                fail("uv.override_package() for '{}': must specify either `target` for full replacement or at least one modification attribute (testonly, console_scripts, pre_build_patches, post_install_patches, exclude_glob, extra_deps, extra_data, toolchains, env, monitor_memory, resource_set).".format(override.name))
 
         unscoped_matches = {i: 0 for i, override in enumerate(mod.tags.override_package) if override.lock == None}
 
@@ -322,6 +324,7 @@ def _parse_projects(module_ctx, hub_specs):
 
             package_overrides = {}
             package_console_scripts = {}
+            testonly_packages = {}
             for i, override in enumerate(mod.tags.override_package):
                 if override.lock != None and override.lock != project.lock:
                     continue
@@ -373,6 +376,8 @@ def _parse_projects(module_ctx, hub_specs):
                 has_target = override.target != None
                 package_overrides[override_key] = override
                 package_console_scripts[override_key] = console_scripts
+                if override.testonly:
+                    testonly_packages[name] = True
 
                 k = (project_id, normalize_name(override.name), v, "__base__")
                 if has_target:
@@ -445,6 +450,23 @@ def _parse_projects(module_ctx, hub_specs):
                     if dep[1] not in marked_package_cfg_sccs:
                         fail("SCC {} depends on package {} without a surface alias".format(scc_id, dep[1]))
 
+            if testonly_packages:
+                for _ in range(len(scc_graph)):
+                    changed = False
+                    for scc_id, members in scc_graph.items():
+                        member_names = [member[1] for member in members]
+                        if not (
+                            any([name in testonly_packages for name in member_names]) or
+                            any([dep[1] in testonly_packages for dep in scc_deps.get(scc_id, {})])
+                        ):
+                            continue
+                        for name in member_names:
+                            if name not in testonly_packages:
+                                testonly_packages[name] = True
+                                changed = True
+                    if not changed:
+                        break
+
             # Pre-build the per-project available_deps mapping from the
             # lockfile. This gives each sdist configure tool visibility
             # into the packages within this project's dependency perimeter.
@@ -617,6 +639,7 @@ def _parse_projects(module_ctx, hub_specs):
                     exclude_glob = exclude_glob,
                     extra_deps = extra_deps,
                     extra_data = extra_data,
+                    testonly = pkg_override.testonly if pkg_override else False,
                 )
 
             # These structures are re-keyed into JSON-serializable shapes for the
@@ -627,6 +650,7 @@ def _parse_projects(module_ctx, hub_specs):
             project_cfgs[project_id] = struct(
                 available_deps = project_available_deps,
                 dep_to_scc = marked_package_cfg_sccs,
+                testonly_packages = testonly_packages,
                 scc_deps = {
                     k: _merge_scc_dep_markers_by_surface_package(deps)
                     for k, deps in scc_deps.items()
@@ -645,7 +669,9 @@ def _parse_projects(module_ctx, hub_specs):
             hub_cfg = hub_cfgs.setdefault(project.hub_name, struct(
                 configurations = {},
                 packages = {},
+                testonly_packages = {},
             ))
+            hub_cfg.testonly_packages.update(testonly_packages)
 
             for cfg in configuration_names.keys():
                 if cfg in hub_cfg.configurations:
@@ -685,6 +711,7 @@ def _parse_projects(module_ctx, hub_specs):
                 },
                 dep_to_scc = pc.dep_to_scc,
                 scc_deps = pc.scc_deps,
+                testonly_packages = pc.testonly_packages,
                 scc_graph = {
                     scc_id: {
                         target_remap.get(target, target): markers
@@ -830,6 +857,8 @@ def _uv_impl(module_ctx):
             install_kwargs["extra_deps"] = json.encode(install_cfg.extra_deps)
         if install_cfg.extra_data:
             install_kwargs["extra_data"] = json.encode(install_cfg.extra_data)
+        if install_cfg.testonly:
+            install_kwargs["package_testonly"] = True
         whl_install(**install_kwargs)
 
     for project_id, project_cfg in cfg.project_cfgs.items():
@@ -839,13 +868,15 @@ def _uv_impl(module_ctx):
             dep_to_scc = json.encode(project_cfg.dep_to_scc),
             scc_deps = json.encode(project_cfg.scc_deps),
             scc_graph = json.encode(project_cfg.scc_graph),
+            testonly_packages = json.encode(project_cfg.testonly_packages),
         )
 
     for hub_id, hub_cfg in cfg.hub_cfgs.items():
         uv_hub(
             name = hub_id,
             configurations = hub_cfg.configurations,
             packages = json.encode(hub_cfg.packages),
+            testonly_packages = json.encode(hub_cfg.testonly_packages),
         )
 
     return module_ctx.extension_metadata(reproducible = True)
@@ -913,6 +944,10 @@ _override_package_tag = tag_class(
             doc = "The `uv.lock` this override applies to. Omit it to apply modifications across every `uv.project()` declared by the same module.",
         ),
         "name": attr.string(mandatory = True),
+        "testonly": attr.bool(
+            default = False,
+            doc = "Restrict this package and its generated Bazel targets to test-only consumers.",
+        ),
         "version": attr.string(mandatory = False),
         "target": attr.label(
             mandatory = False,

uv/private/extension/tests/defs_test.bzl

@@ -78,14 +78,16 @@ def _install_cfg(
         sbuild = None,
         post_install_patches = [],
         extra_deps = [],
-        extra_data = []):
+        extra_data = [],
+        testonly = False):
     return struct(
         whls = whls,
         exclude_glob = exclude_glob,
         sbuild = sbuild,
         post_install_patches = post_install_patches,
         extra_deps = extra_deps,
         extra_data = extra_data,
+        testonly = testonly,
     )
 
 def _shared_install_key_test_impl(ctx):
@@ -115,6 +117,9 @@ def _shared_install_key_test_impl(ctx):
         key == shared_install_key(_install_cfg(ordered, exclude_glob = ["**/*.pyi"])),
     )
 
+    # Test-only installs must never be shared with unrestricted installs.
+    asserts.equals(env, False, key == shared_install_key(_install_cfg(ordered, testonly = True)))
+
     # Project-local inputs are never shareable.
     asserts.equals(env, None, shared_install_key(_install_cfg(ordered, sbuild = "@sdist_build__x//:whl")))
     asserts.equals(env, None, shared_install_key(_install_cfg(ordered, post_install_patches = ["//:demo.patch"])))

uv/private/uv_hub/repository.bzl

@@ -19,6 +19,7 @@ def _hub_impl(repository_ctx):
 
     # {requirement: {cfg: target}}
     packages = json.decode(repository_ctx.attr.packages)
+    testonly_packages = json.decode(repository_ctx.attr.testonly_packages)
     package_names = sorted(packages.keys())
 
     ################################################################################
@@ -84,6 +85,7 @@ exports_files(
     ################################################################################
     # Lay down the hub aliases
     for package_name, specs in packages.items():
+        testonly_attr = "\n    testonly = True," if package_name in testonly_packages else ""
         content = [
             """\
 load("//:defs.bzl", "compatible_with")
@@ -105,11 +107,11 @@ load("//:defs.bzl", "compatible_with")
         # exposes a `pkg` target — emitting a separate `:pkg` alias would collide.
         pkg_alias = "" if package_name == "pkg" else """\
 alias(
-    name = "pkg",
+    name = "pkg",{testonly}
     actual = "{name}",
     visibility = ["//visibility:public"],
 )
-""".format(name = package_name)
+""".format(name = package_name, testonly = testonly_attr)
 
         # FIXME: Add support for entrypoints?
         content.append(
@@ -122,7 +124,7 @@ alias(
     visibility = ["//visibility:public"],
 )
 alias(
-    name = "{name}",
+    name = "{name}",{testonly}
     actual = select({lib_select},
         no_match_error = "{error}",
     ),
@@ -141,6 +143,7 @@ exports_files(
                 whl_select = indent(pprint(whl_select_spec), "      ").lstrip(),
                 compat = repr(specs.keys()),
                 error = error,
+                testonly = testonly_attr,
             ),
         )
 
@@ -304,5 +307,6 @@ uv_hub = repository_rule(
             JSON blob mapping packages to configurations to projects.
             """,
         ),
+        "testonly_packages": attr.string(default = "{}"),
     },
 )

uv/private/uv_project/repository.bzl

@@ -27,6 +27,14 @@ def _project_impl(repository_ctx):
     dep_to_scc = json.decode(repository_ctx.attr.dep_to_scc)
     scc_deps = json.decode(repository_ctx.attr.scc_deps)
     scc_graph = json.decode(repository_ctx.attr.scc_graph)
+    testonly_packages = json.decode(repository_ctx.attr.testonly_packages)
+    testonly_sccs = {
+        scc: True
+        for package, cfgs in dep_to_scc.items()
+        if package in testonly_packages
+        for sccs in cfgs.values()
+        for scc in sccs
+    }
 
     # Collect all the underlying whl installs
     installs = {}
@@ -56,7 +64,7 @@ def _project_impl(repository_ctx):
             acc.append(c if c.isalnum() or c in "._-" else "_")
         return "".join(acc)
 
-    def _conditionalize(it, markers, cond_id_thunk, no_match = None):
+    def _conditionalize(it, markers, cond_id_thunk, no_match = None, testonly = False):
         if "" in markers:
             return it
         else:
@@ -71,13 +79,14 @@ def _project_impl(repository_ctx):
             cond_id = cond_id_thunk()
             content.append("""
 alias(
-    name = "{name}",
+    name = "{name}",{testonly}
     actual = select({arms}),
     visibility = ["//:__subpackages__"],
 )
 """.format(
                 name = cond_id,
                 arms = indent(pprint(cases), " " * 4).lstrip(),
+                testonly = "\n    testonly = True," if testonly else "",
             ))
             return ":" + cond_id
 
@@ -126,6 +135,7 @@ filegroup(
 """.format(package, indent(pprint(cfgs), "# ")))
         main_arms = {}
         whl_main_arms = {}
+        testonly_attr = "\n    testonly = True," if package in testonly_packages else ""
 
         # FIXME: Handle markers for distinct versions
         for cfg, scc_cfgs in cfgs.items():
@@ -143,11 +153,11 @@ filegroup(
 
             content.append("""
 alias(
-    name = "{name}",
+    name = "{name}",{testonly}
     actual = select({arms}),
     visibility = ["//visibility:private"],
 )
-""".format(name = cfg_name, arms = indent(pprint(cfg_arms), " " * 4).lstrip()))
+""".format(name = cfg_name, arms = indent(pprint(cfg_arms), " " * 4).lstrip(), testonly = testonly_attr))
             whl_main_arms["//private/dep_group:" + cfg] = ":" + whl_cfg_name
             content.append("""
 alias(
@@ -159,13 +169,14 @@ alias(
 
         content.append("""
 alias(
-    name = "{name}",
+    name = "{name}",{testonly}
     actual = select({arms}),
     visibility = ["//visibility:public"],
 )
 """.format(
             name = package,
             arms = indent(pprint(main_arms), " " * 4).lstrip(),
+            testonly = testonly_attr,
         ))
 
         content.append("""
@@ -231,6 +242,7 @@ py_library(
                 markers,
                 lambda: "_maybe__{}__{}".format(scc_id, _safe_name(member)),
                 no_match = ":empty",
+                testonly = scc_id in testonly_sccs,
             ))
 
         for dep, markers in this_scc_deps.items():
@@ -240,17 +252,19 @@ py_library(
                 markers,
                 lambda: "_maybe__{}__{}".format(scc_id, _safe_name(dep)),
                 no_match = ":empty",
+                testonly = scc_id in testonly_sccs,
             ))
 
         content.append("""
 py_library(
-    name = "{name}",
+    name = "{name}",{testonly}
     deps = {deps},
     visibility = ["//:__subpackages__"],
 )
 """.format(
             name = scc_id,
             deps = indent(pprint(deps), " " * 4).lstrip(),
+            testonly = "\n    testonly = True," if scc_id in testonly_sccs else "",
         ))
 
     content.append("""
@@ -296,5 +310,6 @@ uv_project = repository_rule(
         "dep_to_scc": attr.string(),
         "scc_deps": attr.string(),
         "scc_graph": attr.string(),
+        "testonly_packages": attr.string(default = "{}"),
     },
 )