Allow specifying the version of buildtools in bzlmod

Resolve keith#86

.bazelrc

@@ -6,8 +6,8 @@ common --repo_env=BAZEL_USE_CPP_ONLY_TOOLCHAIN=1
 # This lets us glob() up all the files inside the examples to make them inputs to tests
 # (Note, we cannot use `common --deleted_packages` because the bazel version command doesn't support it)
 # To update these lines, run `bazel run @rules_bazel_integration_test//tools:update_deleted_packages`.
-build --deleted_packages=examples/bzlmod,examples/simple,examples/specify_assets
-query --deleted_packages=examples/bzlmod,examples/simple,examples/specify_assets
+build --deleted_packages=examples/bzlmod,examples/bzlmod_specify_assets,examples/simple,examples/specify_assets
+query --deleted_packages=examples/bzlmod,examples/bzlmod_specify_assets,examples/simple,examples/specify_assets
 
 # Verbose Failures
 build --verbose_failures

MODULE.bazel

@@ -7,9 +7,10 @@ module(
 bazel_dep(name = "bazel_skylib", version = "1.3.0")
 bazel_dep(name = "platforms", version = "0.0.5")
 
-non_module_deps = use_extension("//:defs.bzl", "buildifier_prebuilt_deps_extension")
+buildifier_prebuilt = use_extension("//:defs.bzl", "buildifier_prebuilt_deps_extension")
+buildifier_prebuilt.toolchains()
 use_repo(
-    non_module_deps,
+    buildifier_prebuilt,
     "buildifier_prebuilt_toolchains",
 )
 

README.md

@@ -98,6 +98,16 @@ the version.
 
 ```bzl
 # Use buildtools version 4.2.5.
+
+# Bzlmod
+buildifier_prebuilt = use_extension("//:defs.bzl", "buildifier_prebuilt_deps_extension")
+buildifier_prebuilt.toolchains(version = "4.2.5")
+use_repo(
+    buildifier_prebuilt,
+    "buildifier_prebuilt_toolchains",
+)
+
+# Workspace
 buildifier_prebuilt_register_toolchains(
     assets = buildtools_assets(version = "4.2.5"),
 )
@@ -114,6 +124,28 @@ To add SHA256 values to the declaration, add a `sha256_values` attribute and spe
 
 ```bzl
 # Use buildtools version 4.2.5.
+
+# Bzlmod
+buildifier_prebuilt = use_extension("//:defs.bzl", "buildifier_prebuilt_deps_extension")
+buildifier_prebuilt.toolchains(
+    sha256_values = {
+        "buildifier_darwin_amd64": "757f246040aceb2c9550d02ef5d1f22d3ef1ff53405fe76ef4c6239ef1ea2cc1",
+        "buildifier_darwin_arm64": "4cf02e051f6cda18765935cb6e77cc938cf8b405064589a50fe9582f82c7edaf",
+        "buildifier_linux_amd64": "f94e71b22925aff76ce01a49e1c6c6d31f521bbbccff047b81f2ea01fd01a945",
+        "buildifier_linux_arm64": "2113d79e45efb51e2b3013c8737cb66cadae3fd89bd7e820438cb06201e50874",
+        "buildozer_darwin_amd64": "3fe671620e6cb7d2386f9da09c1de8de88b02b9dd9275cdecd8b9e417f74df1b",
+        "buildozer_darwin_arm64": "ff4d297023fe3e0fd14113c78f04cef55289ca5bfe5e45a916be738b948dc743",
+        "buildozer_linux_amd64": "e8e39b71c52318a9030dd9fcb9bbfd968d0e03e59268c60b489e6e6fc1595d7b",
+        "buildozer_linux_arm64": "96227142969540def1d23a9e8225524173390d23f3d7fd56ce9c4436953f02fc",
+    },
+    version = "4.2.5",
+)
+use_repo(
+    buildifier_prebuilt,
+    "buildifier_prebuilt_toolchains",
+)
+
+# Workspace
 buildifier_prebuilt_register_toolchains(
     assets = buildtools_assets(version = "4.2.5"),
     sha256_values = {

defs.bzl

@@ -63,7 +63,7 @@ def buildifier_prebuilt_register_toolchains(
                     name = asset.name,
                     platform = asset.platform,
                     arch = asset.arch,
-                    ext = http_file_ext
+                    ext = http_file_ext,
                 ),
             ],
             "downloaded_file_path": asset.name + http_file_ext,
@@ -89,7 +89,46 @@ def buildifier_prebuilt_register_toolchains(
 buildtools_asset = buildtools.create_asset
 buildtools_assets = buildtools.create_assets
 
-def _impl(_):
-    buildifier_prebuilt_register_toolchains(register_toolchains = False)
+def _buildifier_prebuilt_deps_extension_impl(module_ctx):
+    registrations = {}
+    for mod in module_ctx.modules:
+        for toolchains in mod.tags.toolchains:
+            if toolchains.name in registrations:
+                continue
+            registrations[toolchains.name] = True
+            if not toolchains.version:
+                assets = buildtools.DEFAULT_ASSETS
+            else:
+                kwargs = {}
+                if toolchains.names:
+                    kwargs["names"] = toolchains.names
+                if toolchains.platforms:
+                    kwargs["platforms"] = toolchains.platforms
+                if toolchains.arches:
+                    kwargs["arches"] = toolchains.arches
+                if toolchains.sha256_values:
+                    kwargs["sha256_values"] = toolchains.sha256_values
+                assets = buildtools_assets(
+                    version = toolchains.version,
+                    **kwargs
+                )
+            buildifier_prebuilt_register_toolchains(
+                name = toolchains.name,
+                assets = assets,
+                register_toolchains = toolchains.register,
+            )
 
-buildifier_prebuilt_deps_extension = module_extension(implementation = _impl)
+buildifier_prebuilt_deps_extension = module_extension(
+    implementation = _buildifier_prebuilt_deps_extension_impl,
+    tag_classes = {
+        "toolchains": tag_class(attrs = {
+            "name": attr.string(doc = "The name of the repository used for the toolchains.", default = "buildifier_prebuilt_toolchains"),
+            "version": attr.string(doc = "The buildtools version string."),
+            "names": attr.string_list(doc = "A list of tools to include."),
+            "platforms": attr.string_list(doc = "A list of platforms to include."),
+            "arches": attr.string_list(doc = "A list of arches to include."),
+            "sha256_values": attr.string_dict(doc = "A dict of asset name to sha256."),
+            "register": attr.bool(doc = "Whether to register the toolchains."),
+        }),
+    },
+)

examples/BUILD

@@ -33,6 +33,7 @@ sh_binary(
 
 _ALL_EXAMPLES = [
     "bzlmod",
+    "bzlmod_specify_assets",
     "simple",
     "specify_assets",
 ]
@@ -44,6 +45,7 @@ _EXAMPLE_RUNNERS = {
 
 _EXAMPLE_BAZEL_VERSIONS = {
     "bzlmod": [bazel_binaries.versions.current],
+    "bzlmod_specify_assets": [bazel_binaries.versions.current],
 }
 
 [

examples/bzlmod_specify_assets/.bazelrc

@@ -0,0 +1,2 @@
+# Enable bzlmod
+common --enable_bzlmod

examples/bzlmod_specify_assets/BUILD

@@ -0,0 +1,29 @@
+load("@buildifier_prebuilt//:rules.bzl", "buildifier", "buildifier_test")
+
+buildifier(
+    name = "buildifier.check",
+    disabled_rewrites = [
+        "label",
+    ],
+    exclude_patterns = [
+        "./.git/*",
+    ],
+    lint_mode = "warn",
+)
+
+buildifier_test(
+    name = "buildifier.test",
+    srcs = glob(["**/BUILD"]),
+    lint_mode = "warn",
+)
+
+sh_test(
+    name = "version_test",
+    srcs = ["version_test.sh"],
+    args = [
+        "$(location @buildifier_prebuilt//:buildozer)",
+    ],
+    data = [
+        "@buildifier_prebuilt//:buildozer",
+    ],
+)

examples/bzlmod_specify_assets/MODULE.bazel

@@ -0,0 +1,29 @@
+module(
+    name = "bzlmod_example",
+    version = "0.0.0",
+)
+
+bazel_dep(name = "buildifier_prebuilt", version = "0.0.0")
+local_path_override(
+    module_name = "buildifier_prebuilt",
+    path = "../..",
+)
+
+buildifier_prebuilt = use_extension("@buildifier_prebuilt//:defs.bzl", "buildifier_prebuilt_deps_extension")
+buildifier_prebuilt.toolchains(
+    sha256_values = {
+        "buildifier_darwin_amd64": "757f246040aceb2c9550d02ef5d1f22d3ef1ff53405fe76ef4c6239ef1ea2cc1",
+        "buildifier_darwin_arm64": "4cf02e051f6cda18765935cb6e77cc938cf8b405064589a50fe9582f82c7edaf",
+        "buildifier_linux_amd64": "f94e71b22925aff76ce01a49e1c6c6d31f521bbbccff047b81f2ea01fd01a945",
+        "buildifier_linux_arm64": "2113d79e45efb51e2b3013c8737cb66cadae3fd89bd7e820438cb06201e50874",
+        "buildozer_darwin_amd64": "3fe671620e6cb7d2386f9da09c1de8de88b02b9dd9275cdecd8b9e417f74df1b",
+        "buildozer_darwin_arm64": "ff4d297023fe3e0fd14113c78f04cef55289ca5bfe5e45a916be738b948dc743",
+        "buildozer_linux_amd64": "e8e39b71c52318a9030dd9fcb9bbfd968d0e03e59268c60b489e6e6fc1595d7b",
+        "buildozer_linux_arm64": "96227142969540def1d23a9e8225524173390d23f3d7fd56ce9c4436953f02fc",
+    },
+    version = "4.2.5",
+)
+use_repo(
+    buildifier_prebuilt,
+    "buildifier_prebuilt_toolchains",
+)

examples/bzlmod_specify_assets/WORKSPACE

@@ -0,0 +1 @@
+# Intentionally blank. Mark the root of the workspace.

examples/bzlmod_specify_assets/version_test.sh

@@ -0,0 +1,6 @@
+#!/usr/bin/env bash
+
+if "$1" --version | grep -q "4.2.5"; then
+  exit 0
+fi
+exit 1