Init toolchain

Author: groodt

MODULE.bazel

@@ -68,20 +68,11 @@ pip.parse(
 )
 use_repo(pip, "rules_python_publish_deps")
 
-uv = use_extension("//python/private/uv:uv.bzl", "uv")
-uv.download(
-    files = {
-        "uv-aarch64-apple-darwin.tar.gz": "f588388d2b13f77e4526e619f618a306b6b026a96975fbfb2c6dd1ded134cb72",
-        "uv-aarch64-unknown-linux-gnu.tar.gz": "f342442088a56a8a5e4af6781501870bed1b388b37ac2e9deb250cd1d0dc1845",
-        "uv-powerpc64le-unknown-linux-gnu.tar.gz": "0f38a41264be0ef325f8d438f34ea95c002736b3c092e3276518a2253c4ff923",
-        "uv-s390x-unknown-linux-gnu.tar.gz": "dc843c32e51ee3fb46381699bffbc4b6040c51083a1dc7d6c597582b84956aa7",
-        "uv-x86_64-apple-darwin.tar.gz": "6262eba42ebb9035a574b74c5ea253a41353fb4b6b264600e3b05b1a7f4cabc0",
-        "uv-x86_64-pc-windows-msvc.zip": "898ce45b3767ea97429426a7bd0ec54d070eb9e29ff2b072bdcf288bf6ddc9a1",
-        "uv-x86_64-unknown-linux-gnu.tar.gz": "3f96048fa1b82eca14d45bbcc86714cd0dee19a73ef9311da6707faa308ec25f",
-    },
-    version = "0.2.13",
-)
-use_repo(uv, "rules_python_uv")
+uv = use_extension("//python:extensions.bzl", "uv")
+uv.toolchain(uv_version = "0.2.13")
+use_repo(uv, "uv_toolchains")
+
+register_toolchains("@uv_toolchains//:all")
 
 # ===== DEV ONLY DEPS AND SETUP BELOW HERE =====
 bazel_dep(name = "stardoc", version = "0.6.2", dev_dependency = True, repo_name = "io_bazel_stardoc")

python/BUILD.bazel

@@ -222,6 +222,12 @@ bzl_library(
     ],
 )
 
+bzl_library(
+    name = "toolchain",
+    srcs = [":toolchain.bzl"],
+    visibility = ["//visibility:public"],
+)
+
 bzl_library(
     name = "versions_bzl",
     srcs = ["versions.bzl"],
@@ -336,3 +342,16 @@ exports_files([
 current_py_toolchain(
     name = "current_py_toolchain",
 )
+
+toolchain_type(
+    name = "uv_toolchain_type",
+    visibility = ["//visibility:public"],
+)
+
+resolved_toolchain(
+    name = "uv_resolved_toolchain",
+    # Marked manual so that `bazel test //...` passes
+    # even if no toolchain is registered.
+    tags = ["manual"],
+    visibility = ["//visibility:public"],
+)

python/extensions.bzl

@@ -0,0 +1,46 @@
+"""Extensions for bzlmod."""
+
+load(":repositories.bzl", "uv_register_toolchains")
+
+_DEFAULT_NAME = "uv"
+
+uv_toolchain = tag_class(attrs = {
+    "name": attr.string(doc = """\
+Base name for generated repositories, allowing more than one uv toolchain to be registered.
+Overriding the default is only permitted in the root module.
+""", default = "uv"),
+    "uv_version": attr.string(doc = "Explicit version of uv.", mandatory = True),
+})
+
+def _uv_toolchain_extension(module_ctx):
+    registrations = {}
+    for mod in module_ctx.modules:
+        for toolchain in mod.tags.uv_toolchain:
+            if toolchain.name != _DEFAULT_NAME and not mod.is_root:
+                fail("""\
+                Only the root module may override the default name for the uv toolchain.
+                This prevents conflicting registrations in the global namespace of external repos.
+                """)
+            if toolchain.name not in registrations.keys():
+                registrations[toolchain.name] = []
+            registrations[toolchain.name].append(toolchain.uv_version)
+    for name, versions in registrations.items():
+        if len(versions) > 1:
+            # TODO: should be semver-aware, using MVS
+            selected = sorted(versions, reverse = True)[0]
+
+            # buildifier: disable=print
+            print("NOTE: uv toolchain {} has multiple versions {}, selected {}".format(name, versions, selected))
+        else:
+            selected = versions[0]
+
+        uv_register_toolchains(
+            name = name,
+            uv_version = selected,
+            register = False,
+        )
+
+uv = module_extension(
+    implementation = _uv_toolchain_extension,
+    tag_classes = {"toolchain": uv_toolchain},
+)

python/private/pypi/BUILD.bazel

@@ -236,3 +236,21 @@ bzl_library(
     srcs = ["whl_target_platforms.bzl"],
     deps = [":parse_whl_name_bzl"],
 )
+
+bzl_library(
+    name = "toolchains_repo",
+    srcs = ["toolchains_repo.bzl"],
+    visibility = ["//python:__subpackages__"],
+)
+
+bzl_library(
+    name = "versions",
+    srcs = ["versions.bzl"],
+    visibility = ["//python:__subpackages__"],
+)
+
+bzl_library(
+    name = "resolved_toolchain",
+    srcs = ["resolved_toolchain.bzl"],
+    visibility = ["//python:__subpackages__"],
+)

python/private/pypi/resolved_toolchain.bzl

@@ -0,0 +1,41 @@
+# Copyright 2024 The Bazel Authors. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""This module implements an alias rule to the resolved toolchain.
+"""
+
+DOC = """\
+Exposes a concrete toolchain which is the result of Bazel resolving the
+toolchain for the execution or target platform.
+Workaround for https://github.com/bazelbuild/bazel/issues/14009
+"""
+
+# Forward all the providers
+def _uv_resolved_toolchain_impl(ctx):
+    toolchain_info = ctx.toolchains["//python:uv_toolchain_type"]
+    return [
+        toolchain_info,
+        toolchain_info.defaultinfo,
+        toolchain_info.templatevariableinfo,
+        toolchain_info.uvtoolchaininfo,
+    ]
+
+# Copied from java_toolchain_alias
+# https://cs.opensource.google/bazel/bazel/+/master:tools/jdk/java_toolchain_alias.bzl
+uv_resolved_toolchain = rule(
+    implementation = _uv_resolved_toolchain_impl,
+    toolchains = ["//python:uv_toolchain_type"],
+    incompatible_use_toolchain_transition = True,
+    doc = DOC,
+)

python/private/pypi/toolchains_repo.bzl

@@ -0,0 +1,78 @@
+# Copyright 2024 The Bazel Authors. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"Creates a repository to hold toolchains"
+
+UV_PLATFORMS = {
+    "aarch64-apple-darwin": struct(
+        compatible_with = [
+            "@platforms//os:macos",
+            "@platforms//cpu:aarch64",
+        ],
+    ),
+    "x86_64-apple-darwin": struct(
+        compatible_with = [
+            "@platforms//os:macos",
+            "@platforms//cpu:x86_64",
+        ],
+    ),
+    "x86_64-pc-windows-msvc": struct(
+        compatible_with = [
+            "@platforms//os:windows",
+            "@platforms//cpu:x86_64",
+        ],
+    ),
+    "x86_64-unknown-linux-gnu": struct(
+        compatible_with = [
+            "@platforms//os:linux",
+            "@platforms//cpu:x86_64",
+        ],
+    ),
+}
+
+def _toolchains_repo_impl(repository_ctx):
+    build_content = """# Generated by toolchains_repo.bzl
+#
+# These can be registered in the workspace file or passed to --extra_toolchains flag.
+# By default all these toolchains are registered by the uv_register_toolchains macro
+# so you don't normally need to interact with these targets.
+
+"""
+
+    for [platform, meta] in UV_PLATFORMS.items():
+        build_content += """
+# Declare a toolchain Bazel will select for running the tool in an action
+# on the execution platform.
+toolchain(
+    name = "{platform}_uv_toolchain",
+    exec_compatible_with = {compatible_with},
+    toolchain = "@{user_repository_name}_{platform}//:uv_toolchain",
+    toolchain_type = "@rules_python//python:uv_toolchain_type",
+)
+""".format(
+            platform = platform,
+            user_repository_name = repository_ctx.attr.user_repository_name,
+            compatible_with = meta.compatible_with,
+        )
+
+    repository_ctx.file("BUILD.bazel", build_content)
+
+uv_toolchains_repo = repository_rule(
+    _toolchains_repo_impl,
+    doc = """Creates a repository with toolchain definitions for all known platforms
+     which can be registered or selected.""",
+    attrs = {
+        "user_repository_name": attr.string(doc = "what the user chose for the base name"),
+    },
+)

python/private/pypi/uv_toolchain.bzl

@@ -0,0 +1,88 @@
+# Copyright 2024 The Bazel Authors. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""This module implements the uv toolchain rule"""
+
+UvToolchainInfo = provider(
+    doc = "Information about how to invoke the uv executable.",
+    fields = {
+        "uv_files": """Files required in runfiles to make the uv executable available.
+
+May be empty if the uv_path points to a locally installed uv binary.""",
+        "uv_path": "Path to the uv executable.",
+    },
+)
+
+# Avoid using non-normalized paths (workspace/../other_workspace/path)
+def _to_manifest_path(ctx, file):
+    if file.short_path.startswith("../"):
+        return "external/" + file.short_path[3:]
+    else:
+        return ctx.workspace_name + "/" + file.short_path
+
+def _uv_toolchain_impl(ctx):
+    if ctx.attr.uv_tool and ctx.attr.uv_path:
+        fail("Can only set one of uv_tool or uv_path but both were set.")
+    if not ctx.attr.uv_tool and not ctx.attr.uv_path:
+        fail("Must set one of uv_tool or uv_path.")
+
+    uv_files = []
+    uv_path = ctx.attr.uv_path
+
+    if ctx.attr.uv_tool:
+        uv_files = ctx.attr.uv_tool.files.to_list()
+        uv_path = _to_manifest_path(ctx, tool_files[0])
+
+    # Make the $(UV_BIN) variable available in places like genrules.
+    # See https://docs.bazel.build/versions/main/be/make-variables.html#custom_variables
+    templatevariableinfo = platform_common.TemplateVariableInfo({
+        "UV_BIN": uv_path,
+    })
+    defaultinfo = DefaultInfo(
+        files = depset(uv_files),
+        runfiles = ctx.runfiles(files = uv_files),
+    )
+    uvtoolchaininfo = UvToolchainInfo(
+        uv_path = uv_path,
+        uv_files = uv_files,
+    )
+
+    # Export all the providers inside our ToolchainInfo
+    # so the resolved_toolchain rule can grab and re-export them.
+    toolchaininfo = platform_common.ToolchainInfo(
+        uvtoolchaininfo = uvtoolchaininfo,
+        templatevariableinfo = templatevariableinfo,
+        defaultinfo = defaultinfo,
+    )
+    return [
+        defaultinfo,
+        toolchaininfo,
+        templatevariableinfo,
+    ]
+
+uv_toolchain = rule(
+    implementation = _uv_toolchain_impl,
+    attrs = {
+        "uv_path": attr.string(
+            doc = "Path to an existing uv executable",
+            mandatory = False,
+        ),
+        "uv_tool": attr.label(
+            doc = "A hermetically downloaded executable target for the target platform.",
+            mandatory = False,
+            allow_single_file = True,
+        ),
+    },
+    doc = "Defines a uv toolchain.",
+)

python/private/pypi/versions.bzl

@@ -0,0 +1,14 @@
+"""Version and integrity information for downloaded artifacts"""
+
+# From: https://github.com/astral-sh/uv/releases
+UV_TOOL_VERSIONS = {
+    "0.2.13": {
+        "aarch64-apple-darwin": "f588388d2b13f77e4526e619f618a306b6b026a96975fbfb2c6dd1ded134cb72",
+        "aarch64-unknown-linux-gnu": "f342442088a56a8a5e4af6781501870bed1b388b37ac2e9deb250cd1d0dc1845",
+        "powerpc64le-unknown-linux-gnu": "0f38a41264be0ef325f8d438f34ea95c002736b3c092e3276518a2253c4ff923",
+        "s390x-unknown-linux-gnu": "dc843c32e51ee3fb46381699bffbc4b6040c51083a1dc7d6c597582b84956aa7",
+        "x86_64-apple-darwin": "6262eba42ebb9035a574b74c5ea253a41353fb4b6b264600e3b05b1a7f4cabc0",
+        "x86_64-pc-windows-msvc": "898ce45b3767ea97429426a7bd0ec54d070eb9e29ff2b072bdcf288bf6ddc9a1",
+        "x86_64-unknown-linux-gnu": "3f96048fa1b82eca14d45bbcc86714cd0dee19a73ef9311da6707faa308ec25f",
+    },
+}