reimplement the expand_template in a custom rule to better handle selects in the output attr

Author: aignas

python/uv/private/BUILD.bazel

@@ -43,9 +43,10 @@ bzl_library(
         ":toolchain_types_bzl",
         "//python:py_binary_bzl",
         "//python/private:bzlmod_enabled_bzl",
+        "//python/private:full_version_bzl",
         "//python/private:toolchain_types_bzl",
         "@bazel_skylib//lib:shell",
-        "@bazel_skylib//rules:expand_template",
+        "@pythons_hub//:versions_bzl",
     ],
 )
 

python/uv/private/lock.bzl

@@ -16,14 +16,17 @@
 """
 
 load("@bazel_skylib//lib:shell.bzl", "shell")
-load("@bazel_skylib//rules:expand_template.bzl", "expand_template")
+load("@pythons_hub//:versions.bzl", "DEFAULT_PYTHON_VERSION", "MINOR_MAPPING")
 load("//python:py_binary.bzl", "py_binary")
 load("//python/private:bzlmod_enabled.bzl", "BZLMOD_ENABLED")  # buildifier: disable=bzl-visibility
+load("//python/private:full_version.bzl", "full_version")
 load("//python/private:toolchain_types.bzl", "TARGET_TOOLCHAIN_TYPE")  # buildifier: disable=bzl-visibility
 load(":toolchain_types.bzl", "UV_TOOLCHAIN_TYPE")
 
 visibility(["//..."])
 
+_PYTHON_VERSION_FLAG = "//python/config_settings:python_version"
+
 _RunLockInfo = provider(
     doc = "",
     fields = {
@@ -72,7 +75,10 @@ def _args(ctx):
 
 def _lock_impl(ctx):
     srcs = ctx.files.srcs
-    output = ctx.actions.declare_file(ctx.label.name + ".out")
+    output = ctx.actions.declare_file("{}.{}.out".format(
+        ctx.label.name,
+        ctx.attr.python_version.replace(".", "_"),
+    ))
 
     toolchain_info = ctx.toolchains[UV_TOOLCHAIN_TYPE]
     uv = toolchain_info.uv_toolchain_info.uv[DefaultInfo].files_to_run.executable
@@ -87,10 +93,9 @@ def _lock_impl(ctx):
         "--no-python-downloads",
         "--no-cache",
     ])
-    args.add("--custom-compile-command", "bazel run //{}:{}".format(
-        ctx.label.package,
-        ctx.attr.update_target,
-    ))
+    pkg = ctx.label.package
+    update_target = ctx.attr.update_target
+    args.add("--custom-compile-command", "bazel run //{}:{}".format(pkg, update_target))
     if ctx.attr.generate_hashes:
         args.add("--generate-hashes")
     if not ctx.attr.strip_extras:
@@ -151,6 +156,28 @@ def _lock_impl(ctx):
         ),
     ]
 
+def _transition_impl(input_settings, attr):
+    settings = {
+        _PYTHON_VERSION_FLAG: input_settings[_PYTHON_VERSION_FLAG],
+    }
+    if attr.python_version:
+        # FIXME @aignas 2025-03-20: using `full_version` is a workaround for a bug in
+        # how we order toolchains in bazel. If I set the `python_version` flag
+        # to `3.12`, I would expect the latest version to be selected, i.e. the
+        # one that is in MINOR_MAPPING, but it seems that 3.12.0 is selected,
+        # because of how the targets are ordered.
+        settings[_PYTHON_VERSION_FLAG] = full_version(
+            version = attr.python_version,
+            minor_mapping = MINOR_MAPPING,
+        )
+    return settings
+
+_python_version_transition = transition(
+    implementation = _transition_impl,
+    inputs = [_PYTHON_VERSION_FLAG],
+    outputs = [_PYTHON_VERSION_FLAG],
+)
+
 _lock = rule(
     implementation = _lock_impl,
     doc = """\
@@ -182,8 +209,16 @@ modifications and the locking is not done from scratch.
             doc = "Public, see the docs in the macro.",
             default = True,
         ),
+        "output": attr.string(
+            doc = "Public, see the docs in the macro.",
+            mandatory = True,
+        ),
         "python_version": attr.string(
             doc = "Public, see the docs in the macro.",
+            default = full_version(
+                version = DEFAULT_PYTHON_VERSION,
+                minor_mapping = MINOR_MAPPING,
+            ),
         ),
         "srcs": attr.label_list(
             mandatory = True,
@@ -200,12 +235,16 @@ modifications and the locking is not done from scratch.
 The string to input for the 'uv pip compile'.
 """,
         ),
+        "_allowlist_function_transition": attr.label(
+            default = "@bazel_tools//tools/allowlists/function_transition_allowlist",
+        ),
     },
     toolchains = [
         UV_TOOLCHAIN_TYPE,
         # FIXME @aignas 2025-03-17: should this be instead EXEC_TOOLCHAIN_TYPE?
         TARGET_TOOLCHAIN_TYPE,
     ],
+    cfg = _python_version_transition,
 )
 
 def _lock_run_impl(ctx):
@@ -293,17 +332,45 @@ def _maybe_file(path):
 
     return None
 
+def _expand_template_impl(ctx):
+    pkg = ctx.label.package
+    update_src = ctx.actions.declare_file(ctx.attr.update_target + ".py")
+    ctx.actions.expand_template(
+        template = ctx.files._template[0],
+        substitutions = {
+            "{{dst}}": "{}/{}".format(pkg, ctx.attr.output),
+            "{{src}}": "{}".format(ctx.files.src[0].short_path),
+            "{{update_target}}": "//{}:{}".format(pkg, ctx.attr.update_target),
+        },
+        output = update_src,
+    )
+    return DefaultInfo(files = depset([update_src]))
+
+_expand_template = rule(
+    implementation = _expand_template_impl,
+    attrs = {
+        "output": attr.string(mandatory = True),
+        "src": attr.label(mandatory = True),
+        "update_target": attr.string(mandatory = True),
+        "_template": attr.label(
+            default = "//python/uv/private:lock_copier.py",
+            allow_single_file = True,
+        ),
+    },
+    doc = "Expand the template for the update script allowing us to use `select` statements in the {attr}`output` attribute.",
+)
+
 def lock(
         *,
         name,
+        srcs,
+        out,
         args = [],
         build_constraints = [],
         constraints = [],
         env = None,
         generate_hashes = True,
-        out,
         python_version = None,
-        srcs,
         strip_extras = False,
         **kwargs):
     """Pin the requirements based on the src files.
@@ -344,7 +411,6 @@ def lock(
     """
     update_target = "{}.update".format(name)
     locker_target = "{}.run".format(name)
-    template_target = "_{}.update_template".format(name)
 
     # Check if the output file already exists, if yes, first copy it to the
     # output file location in order to make `uv` not change the requirements if
@@ -372,6 +438,7 @@ def lock(
         strip_extras = strip_extras,
         target_compatible_with = target_compatible_with,
         update_target = update_target,
+        output = out,
         tags = [
             "local",
             "manual",
@@ -393,28 +460,19 @@ def lock(
         tags = ["manual"],
     )
 
-    # Write a script that can be used for updating the in-tree version of the
-    # requirements file
-    pkg = native.package_name()
-    expand_template(
-        name = template_target,
-        out = update_target + ".py",
-        template = "//python/uv/private:lock_copier.py",
-        substitutions = {
-            "{{dst}}": "{}/{}".format(pkg, out),
-            "{{src}}": "{}/{}.out".format(pkg, name),
-            "{{update_target}}": "//{}:{}".format(pkg, update_target),
-        },
-        tags = ["manual"],
+    # FIXME @aignas 2025-03-20: is it possible to extend `py_binary` so that the
+    # srcs are generated before `py_binary` is run? I found that
+    # `ctx.files.srcs` usage in the base implementation is making it difficult.
+    _expand_template(
+        name = name + "_cp",
+        src = name,
+        output = out,
+        update_target = update_target,
     )
 
-    # TODO @aignas 2025-03-17: use a custom `py_binary` rule and
-    # `ctx.actions.expand_template` instead of the bazel skylib rule. This
-    # would make it possible to avoid the extra target.
     py_binary(
         name = update_target,
-        srcs = [update_target + ".py"],
-        data = [name] + ([] if not maybe_out else [maybe_out]),
-        python_version = python_version,
-        **kwargs
+        srcs = [name + "_cp"],
+        data = [name] + ([maybe_out] if maybe_out else []),
+        tags = ["manual"],
     )

tests/uv/lock/lock_run_test.py

@@ -42,7 +42,7 @@ def test_requirements_updating_for_the_first_time(self):
             # Then
             self.assertEqual(0, output.returncode, output.stderr)
             self.assertIn(
-                "cp <bazel-sandbox>/tests/uv/lock/requirements_new_file.out",
+                "cp <bazel-sandbox>/tests/uv/lock/requirements_new_file",
                 output.stdout.decode("utf-8"),
             )
             self.assertTrue(want_path.exists(), "The path should exist after the test")
@@ -81,7 +81,7 @@ def test_requirements_updating(self):
             # Then
             self.assertEqual(0, output.returncode)
             self.assertIn(
-                "cp <bazel-sandbox>/tests/uv/lock/requirements.out",
+                "cp <bazel-sandbox>/tests/uv/lock/requirements",
                 output.stdout.decode("utf-8"),
             )
             self.assertEqual(want_path.read_text(), want_text)