fix(uv): Preserve uv wrapper runfiles in lock targets

news/4181.fixed.md

@@ -0,0 +1,4 @@
+(`uv`) Support executable wrappers with multiple outputs and preserve their
+runfiles in lock actions and runnable targets. Export downloaded `uv` binaries
+so wrappers can declare them as dependencies.
+([#4181](https://github.com/bazel-contrib/rules_python/issues/4181))

python/uv/private/lock.bzl

@@ -30,7 +30,9 @@ _RunLockInfo = provider(
     fields = {
         "args": "The args passed to the `uv` by default when running the runnable target.",
         "env": "The env passed to the execution.",
-        "srcs": "Source files required to run the runnable target.",
+        # Preserve the wrapper's runtime files and symlink mappings together;
+        # a regular depset cannot represent the full runfiles layout.
+        "runfiles": "Runtime files required by the runnable target.",
         "template": "The template file for writing a script.",
     },
 )
@@ -129,7 +131,8 @@ def _common_lock(ctx, locker):
 
     output = ctx.actions.declare_file(fname)
     toolchain_info = ctx.toolchains[UV_TOOLCHAIN_TYPE]
-    uv = toolchain_info.uv_toolchain_info.uv[DefaultInfo].files_to_run.executable
+    uv_default_info = toolchain_info.uv_toolchain_info.uv[DefaultInfo]
+    uv = uv_default_info.files_to_run.executable
 
     args = _args(ctx)
     args.add(uv)
@@ -259,7 +262,7 @@ def _common_lock(ctx, locker):
         # exec "$@" in the .sh script.
         arguments = [args.run_shell] if not ctx.attr.is_windows else [],
         tools = [
-            uv,
+            uv_default_info.files_to_run,
             python_files,
             script,
         ],
@@ -278,10 +281,10 @@ def _common_lock(ctx, locker):
         _RunLockInfo(
             args = args.run_info,
             env = ctx.attr.env,
-            srcs = depset(
-                srcs + [uv],
-                transitive = [python_files],
-            ),
+            runfiles = ctx.runfiles(
+                files = srcs + [uv],
+                transitive_files = python_files,
+            ).merge(uv_default_info.default_runfiles),
             template = ctx.files._template[0],
         ),
     ]
@@ -501,7 +504,7 @@ def _run_impl(ctx):
     return [
         DefaultInfo(
             executable = executable,
-            runfiles = ctx.runfiles(transitive_files = info.srcs),
+            runfiles = info.runfiles,
         ),
         RunEnvironmentInfo(
             environment = info.env,

python/uv/private/uv_repository.bzl

@@ -24,6 +24,11 @@ UV_BUILD_TMPL = """\
 # Generated by repositories.bzl
 load("@rules_python//python/uv:uv_toolchain.bzl", "uv_toolchain")
 
+# Exposed for users (e.g. custom wrapper scripts), though depending on a
+# specific toolchain repository's binary directly is ill-advised because
+# toolchain resolution may resolve a different target.
+exports_files(["{binary}"], visibility = ["//visibility:public"])
+
 uv_toolchain(
     name = "uv_toolchain",
     uv = "{binary}",

python/uv/private/uv_toolchain.bzl

@@ -49,9 +49,21 @@ uv_toolchain = rule(
     implementation = _uv_toolchain_impl,
     attrs = {
         "uv": attr.label(
-            doc = "A static uv binary.",
+            doc = """
+The `uv` executable or a wrapper that forwards its arguments to `uv`.
+Runtime dependencies belong in the executable target's runfiles.
+
+:::{versionchanged} VERSION_NEXT_PATCH
+Executable targets with multiple output files are supported. Lock actions
+and runnable targets include the executable target's runfiles.
+:::
+""",
             mandatory = True,
-            allow_single_file = True,
+            # allow_files = True is used instead of allow_single_file = True
+            # because executable rules like py_binary include additional files
+            # (e.g. srcs) in DefaultInfo.files. executable = True still enforces
+            # a single executable.
+            allow_files = True,
             executable = True,
             cfg = "exec",
         ),

tests/uv/lock/lock_tests.bzl

@@ -17,12 +17,38 @@
 load("@bazel_skylib//rules:diff_test.bzl", "diff_test")
 load("@bazel_skylib//rules:native_binary.bzl", "native_test")
 load("@rules_testing//lib:test_suite.bzl", "test_suite")
+load("//python:py_binary.bzl", "py_binary")
 load("//python/uv:lock.bzl", "lock")
+load("//python/uv:uv_toolchain.bzl", "uv_toolchain")
 load("//python/uv/private:lock.bzl", lock_testing = "testing")  # buildifier: disable=bzl-visibility
 load("//tests/support:py_reconfig.bzl", "py_reconfig_test")
 
 _basic_tests = []
 
+def _extra_runfiles_impl(ctx):
+    symlink_payload = ctx.actions.declare_file(
+        ctx.label.name + ".symlink_payload",
+    )
+    ctx.actions.write(symlink_payload, "symlink payload\n")
+    root_symlink_payload = ctx.actions.declare_file(
+        ctx.label.name + ".root_symlink_payload",
+    )
+    ctx.actions.write(root_symlink_payload, "root symlink payload\n")
+
+    # Keep the payloads out of ordinary runfiles to require their symlink mappings.
+    return [DefaultInfo(
+        runfiles = ctx.runfiles(
+            symlinks = {"uv_wrapper/symlink_payload.txt": symlink_payload},
+            root_symlinks = {
+                "uv_wrapper/root_symlink_payload.txt": root_symlink_payload,
+            },
+        ),
+    )]
+
+_extra_runfiles = rule(
+    implementation = _extra_runfiles_impl,
+)
+
 def _test_reroot(env):
     reroot = lock_testing.reroot
     env.expect.that_str(
@@ -190,6 +216,53 @@ def lock_test_suite(name):
         }),
     )
 
+    _extra_runfiles(
+        name = "uv_extra_runfiles",
+    )
+
+    py_binary(
+        name = "uv_with_runfiles",
+        srcs = ["uv_with_runfiles.py"],
+        data = [
+            "testdata/toolchain_payload.txt",
+            ":uv_extra_runfiles",
+        ],
+        deps = ["//python/runfiles"],
+    )
+
+    uv_toolchain(
+        name = "uv_with_runfiles_impl",
+        uv = ":uv_with_runfiles",
+        version = "0.0.0",
+    )
+
+    native.toolchain(
+        name = "uv_with_runfiles_toolchain",
+        toolchain = ":uv_with_runfiles_impl",
+        toolchain_type = "//python/uv:uv_toolchain_type",
+    )
+
+    lock(
+        name = "toolchain_requirements",
+        srcs = ["testdata/requirements.in"],
+        out = "toolchain_requirements.txt",
+        directory = None,
+    )
+
+    for mode in ["run", "update"]:
+        py_reconfig_test(
+            name = "toolchain_runfiles_" + mode + "_test",
+            srcs = ["toolchain_runfiles_test.py"],
+            main = "toolchain_runfiles_test.py",
+            args = ["$(rlocationpath :toolchain_requirements." + mode + ")"],
+            data = [":toolchain_requirements." + mode],
+            deps = ["//python/runfiles"],
+            extra_toolchains = [
+                str(Label(":uv_with_runfiles_toolchain")),
+                str(Label("//tests/support/cc_toolchains:all")),
+            ],
+        )
+
     test_suite(
         name = name + "_basic",
         basic_tests = _basic_tests,
@@ -199,6 +272,8 @@ def lock_test_suite(name):
         name = name,
         tests = [
             ":" + name + "_basic",
+            ":toolchain_runfiles_run_test",
+            ":toolchain_runfiles_update_test",
             ":requirements_test",
             ":requirements_directory_test",
             "//tests/uv/lock/pyproject_toml:requirements_test",

tests/uv/lock/testdata/toolchain_payload.txt

@@ -0,0 +1 @@
+custom uv toolchain runfiles

tests/uv/lock/toolchain_runfiles_test.py

@@ -0,0 +1,32 @@
+import os
+import subprocess
+import sys
+import tempfile
+import unittest
+from pathlib import Path
+
+from python import runfiles
+
+LAUNCHER = sys.argv.pop(1)
+
+
+class ToolchainRunfilesTest(unittest.TestCase):
+    def test_lock_with_toolchain_runfiles(self):
+        files = runfiles.CreateOrRaise()
+        launcher = files.root() / LAUNCHER
+        with tempfile.TemporaryDirectory() as directory:
+            output = Path(directory, "tests/uv/lock/toolchain_requirements.txt")
+            output.parent.mkdir(parents=True)
+            env = dict(os.environ, BUILD_WORKSPACE_DIRECTORY=directory)
+            env.update(files.EnvVars())
+            env.pop("TEST_SRCDIR", None)
+            command = [str(launcher)]
+            if os.name == "nt" and launcher.suffix == ".bat":
+                command = ["cmd.exe", "/c", str(launcher)]
+            result = subprocess.run(command, env=env, capture_output=True, text=True)
+            self.assertEqual(result.returncode, 0, result.stdout + result.stderr)
+            self.assertEqual(output.read_text(), "custom uv toolchain runfiles\n")
+
+
+if __name__ == "__main__":
+    unittest.main()

tests/uv/lock/uv_with_runfiles.py

@@ -0,0 +1,22 @@
+import argparse
+from pathlib import Path
+
+from python import runfiles
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--output-file", type=Path, required=True)
+    args, _ = parser.parse_known_args()
+    root = runfiles.CreateOrRaise().root()
+    for location, expected in [
+        ("_main/uv_wrapper/symlink_payload.txt", "symlink payload\n"),
+        ("uv_wrapper/root_symlink_payload.txt", "root symlink payload\n"),
+    ]:
+        assert (root / location).read_text() == expected, location
+    payload = root / "_main/tests/uv/lock/testdata/toolchain_payload.txt"
+    args.output_file.write_bytes(payload.read_bytes())
+
+
+if __name__ == "__main__":
+    main()