feat: stop autoregistering toolchains

MODULE.bazel

@@ -14,17 +14,25 @@ bazel_dep(name = "platforms", version = "0.0.7")
 # 0.5.4 is the first version with bzlmod support
 bazel_dep(name = "stardoc", version = "0.5.4", repo_name = "io_bazel_stardoc")
 
-ext = use_extension("@aspect_bazel_lib//lib:extensions.bzl", "ext")
-use_repo(ext, "copy_directory_toolchains", "copy_to_directory_toolchains", "coreutils_toolchains", "expand_template_toolchains", "jq_toolchains", "tar_toolchains", "yq_toolchains")
+bazel_lib_toolchains = use_extension("@aspect_bazel_lib//lib:extensions.bzl", "toolchains", dev_dependency = True)
+bazel_lib_toolchains.copy_directory()
+bazel_lib_toolchains.copy_to_directory()
+bazel_lib_toolchains.jq()
+bazel_lib_toolchains.yq()
+bazel_lib_toolchains.coreutils()
+bazel_lib_toolchains.tar()
+bazel_lib_toolchains.expand_template()
+use_repo(bazel_lib_toolchains, "bsd_tar_toolchains", "copy_directory_toolchains", "copy_to_directory_toolchains", "coreutils_toolchains", "expand_template_toolchains", "jq_toolchains", "yq_toolchains")
 
 register_toolchains(
     "@copy_directory_toolchains//:all",
     "@copy_to_directory_toolchains//:all",
     "@jq_toolchains//:all",
-    "@tar_toolchains//:all",
     "@yq_toolchains//:all",
+    "@bsd_tar_toolchains//:all",
     "@coreutils_toolchains//:all",
     "@expand_template_toolchains//:all",
+    dev_dependency = True,
 )
 
 # To allow /tools to be built from source

WORKSPACE

@@ -10,10 +10,20 @@ load(":internal_deps.bzl", "bazel_lib_internal_deps")
 # Fetch deps needed only locally for development
 bazel_lib_internal_deps()
 
-load("//lib:repositories.bzl", "aspect_bazel_lib_dependencies")
+load("//lib:repositories.bzl", "aspect_bazel_lib_dependencies", "register_copy_directory_toolchains", "register_copy_to_directory_toolchains", "register_coreutils_toolchains", "register_expand_template_toolchains", "register_tar_toolchains")
 
 aspect_bazel_lib_dependencies()
 
+register_copy_directory_toolchains()
+
+register_copy_to_directory_toolchains()
+
+register_expand_template_toolchains()
+
+register_coreutils_toolchains()
+
+register_tar_toolchains()
+
 # For running our own unit tests
 load("@bazel_skylib//lib:unittest.bzl", "register_unittest_toolchains")
 

e2e/smoke/MODULE.bazel

@@ -10,3 +10,23 @@ local_path_override(
     module_name = "aspect_bazel_lib",
     path = "../..",
 )
+
+bazel_lib_toolchains = use_extension("@aspect_bazel_lib//lib:extensions.bzl", "toolchains")
+bazel_lib_toolchains.copy_directory()
+bazel_lib_toolchains.copy_to_directory()
+bazel_lib_toolchains.jq()
+bazel_lib_toolchains.yq()
+bazel_lib_toolchains.coreutils()
+bazel_lib_toolchains.tar()
+bazel_lib_toolchains.expand_template()
+use_repo(bazel_lib_toolchains, "bsd_tar_toolchains", "copy_directory_toolchains", "copy_to_directory_toolchains", "coreutils_toolchains", "expand_template_toolchains", "jq_toolchains", "yq_toolchains")
+
+register_toolchains(
+    "@copy_directory_toolchains//:all",
+    "@copy_to_directory_toolchains//:all",
+    "@jq_toolchains//:all",
+    "@yq_toolchains//:all",
+    "@bsd_tar_toolchains//:all",
+    "@coreutils_toolchains//:all",
+    "@expand_template_toolchains//:all",
+)

e2e/smoke/WORKSPACE

@@ -3,13 +3,19 @@ local_repository(
     path = "../..",
 )
 
-load("@aspect_bazel_lib//lib:repositories.bzl", "aspect_bazel_lib_dependencies", "register_jq_toolchains", "register_yq_toolchains")
+load("@aspect_bazel_lib//lib:repositories.bzl", "aspect_bazel_lib_dependencies", "register_copy_directory_toolchains", "register_copy_to_directory_toolchains", "register_coreutils_toolchains", "register_expand_template_toolchains", "register_tar_toolchains")
 
 aspect_bazel_lib_dependencies()
 
-register_jq_toolchains()
+register_copy_directory_toolchains()
 
-register_yq_toolchains()
+register_copy_to_directory_toolchains()
+
+register_expand_template_toolchains()
+
+register_coreutils_toolchains()
+
+register_tar_toolchains()
 
 ############################################
 # rules_go is needed to consume tools from sources

lib/extension_utils.bzl

@@ -0,0 +1,5 @@
+"Public API"
+
+load("//lib/private:extension_utils.bzl", _extension_utils = "extension_utils")
+
+extension_utils = _extension_utils

lib/extensions.bzl

@@ -2,6 +2,16 @@
 
 load(
     "@aspect_bazel_lib//lib:repositories.bzl",
+    "DEFAULT_COPY_DIRECTORY_REPOSITORY",
+    "DEFAULT_COPY_TO_DIRECTORY_REPOSITORY",
+    "DEFAULT_COREUTILS_REPOSITORY",
+    "DEFAULT_COREUTILS_VERSION",
+    "DEFAULT_EXPAND_TEMPLATE_REPOSITORY",
+    "DEFAULT_JQ_REPOSITORY",
+    "DEFAULT_JQ_VERSION",
+    "DEFAULT_TAR_REPOSITORY",
+    "DEFAULT_YQ_REPOSITORY",
+    "DEFAULT_YQ_VERSION",
     "register_copy_directory_toolchains",
     "register_copy_to_directory_toolchains",
     "register_coreutils_toolchains",
@@ -10,17 +20,10 @@ load(
     "register_tar_toolchains",
     "register_yq_toolchains",
 )
+load("//lib:extension_utils.bzl", "extension_utils")
 load("//lib/private:host_repo.bzl", "host_repo")
 
-def _toolchain_extension(mctx):
-    register_copy_directory_toolchains(register = False)
-    register_copy_to_directory_toolchains(register = False)
-    register_jq_toolchains(register = False)
-    register_yq_toolchains(register = False)
-    register_coreutils_toolchains(register = False)
-    register_tar_toolchains(register = False)
-    register_expand_template_toolchains(register = False)
-
+def _host_extension_impl(mctx):
     create_host_repo = False
     for module in mctx.modules:
         if len(module.tags.host) > 0:
@@ -29,8 +32,76 @@ def _toolchain_extension(mctx):
     if create_host_repo:
         host_repo(name = "aspect_bazel_lib_host")
 
-# TODO: some way for users to control repo name/version of the tools installed
-ext = module_extension(
-    implementation = _toolchain_extension,
-    tag_classes = {"host": tag_class(attrs = {})},
+host = module_extension(
+    implementation = _host_extension_impl,
+    tag_classes = {
+        "host": tag_class(attrs = {}),
+    },
+)
+
+def _toolchains_extension_impl(mctx):
+    extension_utils.toolchain_repos(
+        mctx = mctx,
+        get_tag = lambda tags: tags.copy_directory,
+        toolchain_name = "copy_directory",
+        repos_fn = lambda name, version: register_copy_directory_toolchains(name = name, register = False),
+        get_version = lambda attr: None,
+    )
+
+    extension_utils.toolchain_repos(
+        mctx = mctx,
+        get_tag = lambda tags: tags.copy_to_directory,
+        toolchain_name = "copy_to_directory",
+        repos_fn = lambda name, version: register_copy_to_directory_toolchains(name = name, register = False),
+        get_version = lambda attr: None,
+    )
+
+    extension_utils.toolchain_repos(
+        mctx = mctx,
+        get_tag = lambda tags: tags.jq,
+        toolchain_name = "jq",
+        repos_fn = lambda name, version: register_jq_toolchains(name = name, version = version, register = False),
+    )
+
+    extension_utils.toolchain_repos(
+        mctx = mctx,
+        get_tag = lambda tags: tags.yq,
+        toolchain_name = "yq",
+        repos_fn = lambda name, version: register_yq_toolchains(name = name, version = version, register = False),
+    )
+
+    extension_utils.toolchain_repos(
+        mctx = mctx,
+        get_tag = lambda tags: tags.coreutils,
+        toolchain_name = "coreutils",
+        repos_fn = lambda name, version: register_coreutils_toolchains(name = name, version = version, register = False),
+    )
+
+    extension_utils.toolchain_repos(
+        mctx = mctx,
+        get_tag = lambda tags: tags.tar,
+        toolchain_name = "tar",
+        repos_fn = lambda name, version: register_tar_toolchains(name = name, register = False),
+        get_version = lambda attr: None,
+    )
+
+    extension_utils.toolchain_repos(
+        mctx = mctx,
+        get_tag = lambda tags: tags.expand_template,
+        toolchain_name = "expand_template",
+        repos_fn = lambda name, version: register_expand_template_toolchains(name = name, register = False),
+        get_version = lambda attr: None,
+    )
+
+toolchains = module_extension(
+    implementation = _toolchains_extension_impl,
+    tag_classes = {
+        "copy_directory": tag_class(attrs = {"name": attr.string(default = DEFAULT_COPY_DIRECTORY_REPOSITORY)}),
+        "copy_to_directory": tag_class(attrs = {"name": attr.string(default = DEFAULT_COPY_TO_DIRECTORY_REPOSITORY)}),
+        "jq": tag_class(attrs = {"name": attr.string(default = DEFAULT_JQ_REPOSITORY), "version": attr.string(default = DEFAULT_JQ_VERSION)}),
+        "yq": tag_class(attrs = {"name": attr.string(default = DEFAULT_YQ_REPOSITORY), "version": attr.string(default = DEFAULT_YQ_VERSION)}),
+        "coreutils": tag_class(attrs = {"name": attr.string(default = DEFAULT_COREUTILS_REPOSITORY), "version": attr.string(default = DEFAULT_COREUTILS_VERSION)}),
+        "tar": tag_class(attrs = {"name": attr.string(default = DEFAULT_TAR_REPOSITORY)}),
+        "expand_template": tag_class(attrs = {"name": attr.string(default = DEFAULT_EXPAND_TEMPLATE_REPOSITORY)}),
+    },
 )

lib/private/extension_utils.bzl

@@ -0,0 +1,69 @@
+"""Utility functions for bzlmod extensions"""
+
+def _toolchain_repos(mctx, get_tag, toolchain_name, repos_fn, default_repository = None, get_name = None, get_version = None):
+    """Create toolchain repositories from bzlmod extension tags.
+
+    For uses of the tag class where the name isn't overridden (the `default_repository` name
+    is used), creates a single toolchain repository for the invocation of the tag closest to
+    the root module. Otherwise, resolution follows bzlmod's depth-first search through the
+    module graph.
+
+    Only a root module may declare a toolchain with a name other than the default repository name.
+
+    Args:
+        mctx: The module context
+        get_tag: A function that takes in `module.tags` and returns the desired tag. For example,
+          `tag: lambda tags: tags.my_tag`. This is required because `my_tag` cannot be accessed
+          as a simple string key from `module.tags`.
+        toolchain_name: Name of the toolchain to use in error messages
+        repos_fn: A function that takes (name, version) and creates a toolchain repository
+        default_repository: Default name of the toolchain repository to pass to the repos_fn.
+          By default, uses `toolchain_name`.
+        get_name: A function that extracts the module name from the a tag's attributes. Defaults
+          to grabbing the `name` attribute.
+        get_version: A function that extracts the module name from the a tag's attributes. Defaults
+          to grabbing the `version` attribute. Override this to a lambda that returns `None` if
+          version isn't used as an attribute.
+    """
+    if default_repository == None:
+        default_repository = toolchain_name
+
+    if get_name == None:
+        get_name = lambda attr: attr.name
+    if get_version == None:
+        get_version = lambda attr: attr.version
+
+    registrations = {}
+    for mod in mctx.modules:
+        for attr in get_tag(mod.tags):
+            name = get_name(attr)
+            version = get_version(attr)
+            if name != default_repository and not mod.is_root:
+                fail("Only the root module may provide a name for the {}} toolchain.".format(toolchain_name))
+
+            if name in registrations.keys():
+                if name == default_repository:
+                    # Prioritize the root-most registration of the default toolchain version and
+                    # ignore any further registrations (modules are processed breadth-first)
+                    continue
+                if version == registrations[name]:
+                    # No problem to register a matching toolchain twice
+                    continue
+                fail("Multiple conflicting {} toolchains declared for name {} ({} and {})".format(
+                    toolchain_name,
+                    name,
+                    version,
+                    registrations[name],
+                ))
+            else:
+                registrations[name] = version
+
+        for name, version in registrations.items():
+            repos_fn(
+                name = name,
+                version = version,
+            )
+
+extension_utils = struct(
+    toolchain_repos = _toolchain_repos,
+)

lib/repositories.bzl

@@ -23,23 +23,10 @@ def aspect_bazel_lib_dependencies():
         ],
     )
 
-    # Always register the copy_to_directory toolchain
-    register_copy_directory_toolchains()
-    register_copy_to_directory_toolchains()
-
-    # Always register the expand_template toolchain
-    register_expand_template_toolchains()
-
-    # Always register the coreutils toolchain and the tar toolchain
-    register_coreutils_toolchains()
-    register_tar_toolchains()
-
-# Re-export the default versions
+DEFAULT_JQ_REPOSITORY = "jq"
 DEFAULT_JQ_VERSION = _DEFAULT_JQ_VERSION
-DEFAULT_YQ_VERSION = _DEFAULT_YQ_VERSION
-DEFAULT_COREUTILS_VERSION = _DEFAULT_COREUTILS_VERSION
 
-def register_jq_toolchains(name = "jq", version = DEFAULT_JQ_VERSION, register = True):
+def register_jq_toolchains(name = DEFAULT_JQ_REPOSITORY, version = DEFAULT_JQ_VERSION, register = True):
     """Registers jq toolchain and repositories
 
     Args:
@@ -64,7 +51,10 @@ def register_jq_toolchains(name = "jq", version = DEFAULT_JQ_VERSION, register =
         user_repository_name = name,
     )
 
-def register_yq_toolchains(name = "yq", version = DEFAULT_YQ_VERSION, register = True):
+DEFAULT_YQ_REPOSITORY = "yq"
+DEFAULT_YQ_VERSION = _DEFAULT_YQ_VERSION
+
+def register_yq_toolchains(name = DEFAULT_YQ_REPOSITORY, version = DEFAULT_YQ_VERSION, register = True):
     """Registers yq toolchain and repositories
 
     Args:
@@ -89,7 +79,9 @@ def register_yq_toolchains(name = "yq", version = DEFAULT_YQ_VERSION, register =
         user_repository_name = name,
     )
 
-def register_tar_toolchains(name = "bsd_tar", register = True):
+DEFAULT_TAR_REPOSITORY = "bsd_tar"
+
+def register_tar_toolchains(name = DEFAULT_TAR_REPOSITORY, register = True):
     """Registers bsdtar toolchain and repositories
 
     Args:
@@ -110,7 +102,10 @@ def register_tar_toolchains(name = "bsd_tar", register = True):
         user_repository_name = name,
     )
 
-def register_coreutils_toolchains(name = "coreutils", version = DEFAULT_COREUTILS_VERSION, register = True):
+DEFAULT_COREUTILS_REPOSITORY = "coreutils"
+DEFAULT_COREUTILS_VERSION = _DEFAULT_COREUTILS_VERSION
+
+def register_coreutils_toolchains(name = DEFAULT_COREUTILS_REPOSITORY, version = DEFAULT_COREUTILS_VERSION, register = True):
     """Registers coreutils toolchain and repositories
 
     Args:
@@ -133,7 +128,9 @@ def register_coreutils_toolchains(name = "coreutils", version = DEFAULT_COREUTIL
         user_repository_name = name,
     )
 
-def register_copy_directory_toolchains(name = "copy_directory", register = True):
+DEFAULT_COPY_DIRECTORY_REPOSITORY = "copy_directory"
+
+def register_copy_directory_toolchains(name = DEFAULT_COPY_DIRECTORY_REPOSITORY, register = True):
     """Registers copy_directory toolchain and repositories
 
     Args:
@@ -166,7 +163,9 @@ def register_copy_directory_toolchains(name = "copy_directory", register = True)
         user_repository_name = name,
     )
 
-def register_copy_to_directory_toolchains(name = "copy_to_directory", register = True):
+DEFAULT_COPY_TO_DIRECTORY_REPOSITORY = "copy_to_directory"
+
+def register_copy_to_directory_toolchains(name = DEFAULT_COPY_TO_DIRECTORY_REPOSITORY, register = True):
     """Registers copy_to_directory toolchain and repositories
 
     Args:
@@ -199,7 +198,9 @@ def register_copy_to_directory_toolchains(name = "copy_to_directory", register =
         user_repository_name = name,
     )
 
-def register_expand_template_toolchains(name = "expand_template", register = True):
+DEFAULT_EXPAND_TEMPLATE_REPOSITORY = "expand_template"
+
+def register_expand_template_toolchains(name = DEFAULT_EXPAND_TEMPLATE_REPOSITORY, register = True):
     """Registers expand_template toolchain and repositories
 
     Args: