wip

Author: aignas

MODULE.bazel

@@ -173,8 +173,11 @@ use_repo(
     "build_bazel_bazel_self",
 )
 
-# EXPERIMENTAL: This is experimental and may be removed without notice
 uv = use_extension("//python/uv:extensions.bzl", "uv")
+
+# Here is how we can define platforms for the `uv` binaries - this will affect
+# all of the downstream callers because we are using the extension without
+# `dev_dependency = True`.
 uv.platform(
     name = "aarch64-apple-darwin",
     compatible_with = [
@@ -245,10 +248,13 @@ uv_dev = use_extension(
     "uv",
     dev_dependency = True,
 )
-uv_dev.toolchain(version = "0.5.24")
-use_repo(uv_dev, "uv_toolchains")
+uv_dev.toolchain(
+    name = "uv_0_5_24",
+    version = "0.5.24",
+)
+use_repo(uv_dev, "uv_0_5_24")
 
 register_toolchains(
-    "@uv_toolchains//:all",
+    "@uv_0_5_24//:all",
     dev_dependency = True,
 )

examples/bzlmod/MODULE.bazel

@@ -106,7 +106,7 @@ use_repo(python, "python_3_10", "python_3_9", "python_versions", "pythons_hub")
 
 # EXPERIMENTAL: This is experimental and may be removed without notice
 uv = use_extension("@rules_python//python/uv:extensions.bzl", "uv")
-uv.toolchain(uv_version = "0.4.25")
+uv.toolchain(version = "0.4.25")
 use_repo(uv, "uv_toolchains")
 
 register_toolchains("@uv_toolchains//:all")

python/uv/BUILD.bazel

@@ -27,9 +27,6 @@ filegroup(
     visibility = ["//:__subpackages__"],
 )
 
-# For stardoc to reference the files
-exports_files(["defs.bzl"])
-
 toolchain_type(
     name = "uv_toolchain_type",
     visibility = ["//visibility:public"],
@@ -47,13 +44,6 @@ current_toolchain(
     ],
 )
 
-bzl_library(
-    name = "defs",
-    srcs = ["defs.bzl"],
-    # EXPERIMENTAL: Visibility is restricted to allow for changes.
-    visibility = ["//:__subpackages__"],
-)
-
 bzl_library(
     name = "extensions",
     srcs = ["extensions.bzl"],
@@ -74,8 +64,11 @@ bzl_library(
 )
 
 bzl_library(
-    name = "toolchain",
-    srcs = ["toolchain.bzl"],
+    name = "uv_toolchain",
+    srcs = ["uv_toolchain.bzl"],
     # EXPERIMENTAL: Visibility is restricted to allow for changes.
     visibility = ["//:__subpackages__"],
+    deps = [
+        "//python/uv/private:uv_toolchain",
+    ],
 )

python/uv/private/extension.bzl

@@ -48,7 +48,10 @@ platform = tag_class(attrs = {
 })
 
 uv_toolchain = tag_class(attrs = {
-    "uv_version": attr.string(doc = "IGNORED", mandatory = False),
+    "name": attr.string(
+        doc = "The name of the toolchain repo",
+        default = "uv_toolchains",
+    ),
     "version": attr.string(
         doc = "Explicit version of uv.",
         mandatory = True,
@@ -57,21 +60,24 @@ uv_toolchain = tag_class(attrs = {
 
 def _uv_toolchain_extension(module_ctx):
     config = {
-        "base_url": _DEFAULT_BASE_URL,
         "platforms": {},
-        "version": None,
     }
 
     for mod in module_ctx.modules:
-        if not mod.is_root:
-            # Ignore any attempts to configure the `uv` toolchain from non-root modules
+        if not mod.is_root and not mod.name == "rules_python":
+            # Only rules_python and the root module can configure this.
+            #
+            # Ignore any attempts to configure the `uv` toolchain elsewhere
             #
             # Only the root module may configure the uv toolchain.
             # This prevents conflicting registrations with any other modules.
             #
             # NOTE: We may wish to enforce a policy where toolchain configuration is only allowed in the root module, or in rules_python. See https://github.com/bazelbuild/bazel/discussions/22024
             continue
 
+        # Note, that the first registration will always win, givin priority to
+        # the root module.
+
         for platform_attr in mod.tags.platform:
             config["platforms"].setdefault(platform_attr.name, struct(
                 name = platform_attr.name.replace("-", "_").lower(),
@@ -80,19 +86,22 @@ def _uv_toolchain_extension(module_ctx):
             ))
 
         for config_attr in mod.tags.config:
-            config["base_url"] = config_attr.base_url
+            config.setdefault("base_url", config_attr.base_url)
 
         for toolchain in mod.tags.toolchain:
-            config["version"] = toolchain.version
+            config.setdefault("version", toolchain.version)
+            config.setdefault("name", toolchain.name)
 
     if not config["version"]:
         return
 
+    config.setdefault("base_url", _DEFAULT_BASE_URL)
     config["urls"] = _get_tool_urls_from_dist_manifest(
         module_ctx,
         base_url = "{base_url}/{version}".format(**config),
     )
     uv_register_toolchains(
+        name = config["name"],
         platforms = config["platforms"],
         urls = config["urls"],
         version = config["version"],

python/uv/private/repositories.bzl

@@ -76,13 +76,13 @@ uv_repository = repository_rule(
     },
 )
 
-# buildifier: disable=unnamed-macro
-def uv_register_toolchains(*, version, platforms, urls):
+def uv_register_toolchains(*, name, version, platforms, urls):
     """Convenience macro which does typical toolchain setup
 
     Skip this macro if you need more control over the toolchain setup.
 
     Args:
+        name: The name of the toolchains repo,
         version: The uv toolchain version to download.
         platforms: The platforms to register uv for.
         urls: The urls with sha256 values to register uv for.
@@ -94,26 +94,26 @@ def uv_register_toolchains(*, version, platforms, urls):
     toolchain_labels_by_toolchain = {}
     toolchain_compatible_with_by_toolchain = {}
 
-    for platform in platforms.keys():
-        uv_repository_name = platforms[platform].name
+    for platform_name, platform in platforms.items():
+        uv_repository_name = "{}_{}".format(name, platform_name.lower().replace("-", "_"))
         uv_repository(
             name = uv_repository_name,
             version = version,
-            platform = platform,
-            urls = urls[platform].urls,
-            sha256 = urls[platform].sha256,
+            platform = platform_name,
+            urls = urls[platform_name].urls,
+            sha256 = urls[platform_name].sha256,
         )
 
         toolchain_name = uv_repository_name + "_toolchain"
         toolchain_names.append(toolchain_name)
         toolchain_labels_by_toolchain[toolchain_name] = "@{}//:uv_toolchain".format(uv_repository_name)
         toolchain_compatible_with_by_toolchain[toolchain_name] = [
             str(label)
-            for label in platforms[platform].compatible_with
+            for label in platform.compatible_with
         ]
 
     uv_toolchains_repo(
-        name = "uv_toolchains",
+        name = name,
         toolchain_type = str(UV_TOOLCHAIN_TYPE),
         toolchain_names = toolchain_names,
         toolchain_labels = toolchain_labels_by_toolchain,