Add helper functions for module extensions as modules (#456)

Adds a new module `modules` with two helper functions for module
extensions:

* `use_all_repos` makes it easy to return an appropriate
  `extension_metadata` from a module extension (if supported) to
  indicate that all repositories generated by the extension should be
  imported via `use_repo`.
* `as_extension` turns a WORKSPACE macro into a module extension that
  uses `use_all_repos` to automate the generation of `use_repo` calls.

MODULE.bazel

@@ -26,3 +26,9 @@ local_path_override(
     module_name = "bazel_skylib_gazelle_plugin",
     path = "gazelle",
 )
+
+as_extension_test_ext = use_extension("//tests:modules_test.bzl", "as_extension_test_ext")
+use_repo(as_extension_test_ext, "bar", "foo")
+
+use_all_repos_test_ext = use_extension("//tests:modules_test.bzl", "use_all_repos_test_ext")
+use_repo(use_all_repos_test_ext, "baz", "qux")

README.md

@@ -48,6 +48,7 @@ s = shell.quote(p)
 * [paths](docs/paths_doc.md)
 * [selects](docs/selects_doc.md)
 * [sets](lib/sets.bzl) - _deprecated_, use `new_sets`
+* [modules](docs/modules_doc.md)
 * [new_sets](docs/new_sets_doc.md)
 * [shell](docs/shell_doc.md)
 * [structs](docs/structs_doc.md)

docs/BUILD

@@ -62,6 +62,12 @@ stardoc_with_diff_test(
     out_label = "//docs:expand_template_doc.md",
 )
 
+stardoc_with_diff_test(
+    name = "modules",
+    bzl_library_target = "//lib:modules",
+    out_label = "//docs:modules_doc.md",
+)
+
 stardoc_with_diff_test(
     name = "native_binary",
     bzl_library_target = "//rules:native_binary",

docs/modules_doc.md

@@ -0,0 +1,76 @@
+<!-- Generated with Stardoc: http://skydoc.bazel.build -->
+
+Skylib module containing utilities for Bazel modules and module extensions.
+
+<a id="modules.as_extension"></a>
+
+## modules.as_extension
+
+<pre>
+modules.as_extension(<a href="#modules.as_extension-macro">macro</a>, <a href="#modules.as_extension-doc">doc</a>)
+</pre>
+
+Wraps a WORKSPACE dependency macro into a module extension.
+
+Example:
+```starlark
+def rules_foo_deps(optional_arg = True):
+    some_repo_rule(name = "foobar")
+    http_archive(name = "bazqux")
+
+rules_foo_deps_ext = modules.as_extension(rules_foo_deps)
+```
+
+
+**PARAMETERS**
+
+
+| Name  | Description | Default Value |
+| :------------- | :------------- | :------------- |
+| <a id="modules.as_extension-macro"></a>macro |  A [WORKSPACE dependency macro](https://bazel.build/rules/deploying#dependencies), i.e., a function with no required parameters that instantiates one or more repository rules.   |  none |
+| <a id="modules.as_extension-doc"></a>doc |  A description of the module extension that can be extracted by documentation generating tools.   |  <code>None</code> |
+
+**RETURNS**
+
+A module extension that generates the repositories instantiated by the given macro and also
+uses [`use_all_repos`](#use_all_repos) to indicate that all of those repositories should be
+imported via `use_repo`.
+
+
+<a id="modules.use_all_repos"></a>
+
+## modules.use_all_repos
+
+<pre>
+modules.use_all_repos(<a href="#modules.use_all_repos-module_ctx">module_ctx</a>)
+</pre>
+
+Return from a module extension that should have all its repositories imported via `use_repo`.
+
+Example:
+```starlark
+def _ext_impl(module_ctx):
+    some_repo_rule(name = "foobar")
+    http_archive(name = "bazqux")
+    return modules.use_all_repos(module_ctx)
+
+ext = module_extension(_ext_impl)
+```
+
+
+**PARAMETERS**
+
+
+| Name  | Description | Default Value |
+| :------------- | :------------- | :------------- |
+| <a id="modules.use_all_repos-module_ctx"></a>module_ctx |  The [<code>module_ctx</code>](https://bazel.build/rules/lib/builtins/module_ctx) object passed to the module extension's implementation function.   |  none |
+
+**RETURNS**
+
+An [`extension_metadata`](https://bazel.build/rules/lib/builtins/extension_metadata.html)
+object that, when returned from a module extension implementation function, specifies that all
+repositories generated by this extension should be imported via `use_repo`. If the current
+version of Bazel doesn't support `extension_metadata`, returns `None` instead, which can
+safely be returned from a module extension implementation function in all versions of Bazel.
+
+

lib/BUILD

@@ -20,6 +20,11 @@ bzl_library(
     srcs = ["dicts.bzl"],
 )
 
+bzl_library(
+    name = "modules",
+    srcs = ["modules.bzl"],
+)
+
 bzl_library(
     name = "partial",
     srcs = ["partial.bzl"],

lib/modules.bzl

@@ -0,0 +1,99 @@
+# Copyright 2023 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.
+
+"""Skylib module containing utilities for Bazel modules and module extensions."""
+
+def _as_extension(macro, doc = None):
+    """Wraps a WORKSPACE dependency macro into a module extension.
+
+    Example:
+    ```starlark
+    def rules_foo_deps(optional_arg = True):
+        some_repo_rule(name = "foobar")
+        http_archive(name = "bazqux")
+
+    rules_foo_deps_ext = modules.as_extension(rules_foo_deps)
+    ```
+
+    Args:
+      macro: A [WORKSPACE dependency macro](https://bazel.build/rules/deploying#dependencies), i.e.,
+          a function with no required parameters that instantiates one or more repository rules.
+      doc: A description of the module extension that can be extracted by documentation generating
+          tools.
+
+    Returns:
+      A module extension that generates the repositories instantiated by the given macro and also
+      uses [`use_all_repos`](#use_all_repos) to indicate that all of those repositories should be
+      imported via `use_repo`.
+    """
+
+    def _ext_impl(module_ctx):
+        macro()
+        return _use_all_repos(module_ctx)
+
+    return module_extension(
+        implementation = _ext_impl,
+        doc = doc,
+    )
+
+def _use_all_repos(module_ctx):
+    """Return from a module extension that should have all its repositories imported via `use_repo`.
+
+    Example:
+    ```starlark
+    def _ext_impl(module_ctx):
+        some_repo_rule(name = "foobar")
+        http_archive(name = "bazqux")
+        return modules.use_all_repos(module_ctx)
+
+    ext = module_extension(_ext_impl)
+    ```
+
+    Args:
+      module_ctx: The [`module_ctx`](https://bazel.build/rules/lib/builtins/module_ctx) object
+          passed to the module extension's implementation function.
+
+    Returns:
+      An [`extension_metadata`](https://bazel.build/rules/lib/builtins/extension_metadata.html)
+      object that, when returned from a module extension implementation function, specifies that all
+      repositories generated by this extension should be imported via `use_repo`. If the current
+      version of Bazel doesn't support `extension_metadata`, returns `None` instead, which can
+      safely be returned from a module extension implementation function in all versions of Bazel.
+    """
+
+    # module_ctx.extension_metadata is available in Bazel 6.2.0 and later.
+    # If not available, returning None from a module extension is equivalent to not returning
+    # anything.
+    extension_metadata = getattr(module_ctx, "extension_metadata", None)
+    if not extension_metadata:
+        return None
+
+    # module_ctx.root_module_has_non_dev_dependency is available in Bazel 6.3.0 and later.
+    root_module_has_non_dev_dependency = getattr(
+        module_ctx,
+        "root_module_has_non_dev_dependency",
+        None,
+    )
+    if root_module_has_non_dev_dependency == None:
+        return None
+
+    return extension_metadata(
+        root_module_direct_deps = "all" if root_module_has_non_dev_dependency else [],
+        root_module_direct_dev_deps = [] if root_module_has_non_dev_dependency else "all",
+    )
+
+modules = struct(
+    as_extension = _as_extension,
+    use_all_repos = _use_all_repos,
+)

tests/BUILD

@@ -3,6 +3,7 @@ load(":build_test_tests.bzl", "build_test_test_suite")
 load(":collections_tests.bzl", "collections_test_suite")
 load(":common_settings_tests.bzl", "common_settings_test_suite")
 load(":dicts_tests.bzl", "dicts_test_suite")
+load(":modules_test.bzl", "modules_test_suite")
 load(":new_sets_tests.bzl", "new_sets_test_suite")
 load(":partial_tests.bzl", "partial_test_suite")
 load(":paths_tests.bzl", "paths_test_suite")
@@ -29,6 +30,8 @@ common_settings_test_suite()
 
 dicts_test_suite()
 
+modules_test_suite()
+
 new_sets_test_suite()
 
 partial_test_suite()

tests/modules_test.bzl

@@ -0,0 +1,70 @@
+# Copyright 2017 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.
+
+"""Test usage of modules.bzl."""
+
+load("//lib:modules.bzl", "modules")
+load("//rules:build_test.bzl", "build_test")
+
+def _repo_rule_impl(repository_ctx):
+    repository_ctx.file("WORKSPACE")
+    repository_ctx.file("BUILD", """exports_files(["hello"])""")
+    repository_ctx.file("hello", "Hello, Bzlmod!")
+
+_repo_rule = repository_rule(_repo_rule_impl)
+
+def _workspace_macro(register_toolchains = False):
+    _repo_rule(name = "foo")
+    _repo_rule(name = "bar")
+    if register_toolchains:
+        native.register_toolchains()
+
+as_extension_test_ext = modules.as_extension(
+    _workspace_macro,
+    doc = "Only used for testing modules.as_extension().",
+)
+
+def _use_all_repos_ext_impl(module_ctx):
+    _repo_rule(name = "baz")
+    _repo_rule(name = "qux")
+    return modules.use_all_repos(module_ctx)
+
+use_all_repos_test_ext = module_extension(
+    _use_all_repos_ext_impl,
+    doc = "Only used for testing modules.use_all_repos().",
+)
+
+# buildifier: disable=unnamed-macro
+def modules_test_suite():
+    """Creates the tests for modules.bzl if Bzlmod is enabled."""
+
+    is_bzlmod_enabled = str(Label("//tests:module_tests.bzl")).startswith("@@")
+    if not is_bzlmod_enabled:
+        return
+
+    build_test(
+        name = "modules_as_extension_test",
+        targets = [
+            "@foo//:hello",
+            "@bar//:hello",
+        ],
+    )
+
+    build_test(
+        name = "modules_use_all_repos_test",
+        targets = [
+            "@baz//:hello",
+            "@qux//:hello",
+        ],
+    )