refactor(bzlmod): move bzlmod extension implementation to python/priv…

…ate/bzlmod

MODULE.bazel

@@ -14,7 +14,7 @@ bazel_dep(name = "protobuf", version = "21.7", repo_name = "com_google_protobuf"
 
 bazel_dep(name = "stardoc", version = "0.6.2", dev_dependency = True, repo_name = "io_bazel_stardoc")
 
-internal_deps = use_extension("@rules_python//python/extensions/private:internal_deps.bzl", "internal_deps")
+internal_deps = use_extension("@rules_python//python/private/bzlmod:internal_deps.bzl", "internal_deps")
 internal_deps.install()
 use_repo(
     internal_deps,

python/extensions/python.bzl

@@ -14,253 +14,6 @@
 
 "Python toolchain module extensions for use with bzlmod"
 
-load("//python:repositories.bzl", "python_register_toolchains")
-load("//python/extensions/private:pythons_hub.bzl", "hub_repo")
-load("//python/private:toolchains_repo.bzl", "multi_toolchain_aliases")
+load("//python/private/bzlmod:python.bzl", _python = "python")
 
-# This limit can be increased essentially arbitrarily, but doing so will cause a rebuild of all
-# targets using any of these toolchains due to the changed repository name.
-_MAX_NUM_TOOLCHAINS = 9999
-_TOOLCHAIN_INDEX_PAD_LENGTH = len(str(_MAX_NUM_TOOLCHAINS))
-
-def _toolchain_prefix(index, name):
-    """Prefixes the given name with the index, padded with zeros to ensure lexicographic sorting.
-
-    Examples:
-      _toolchain_prefix(   2, "foo") == "_0002_foo_"
-      _toolchain_prefix(2000, "foo") == "_2000_foo_"
-    """
-    return "_{}_{}_".format(_left_pad_zero(index, _TOOLCHAIN_INDEX_PAD_LENGTH), name)
-
-def _left_pad_zero(index, length):
-    if index < 0:
-        fail("index must be non-negative")
-    return ("0" * length + str(index))[-length:]
-
-# Printing a warning msg not debugging, so we have to disable
-# the buildifier check.
-# buildifier: disable=print
-def _print_warn(msg):
-    print("WARNING:", msg)
-
-def _python_register_toolchains(name, toolchain_attr, version_constraint):
-    """Calls python_register_toolchains and returns a struct used to collect the toolchains.
-    """
-    python_register_toolchains(
-        name = name,
-        python_version = toolchain_attr.python_version,
-        register_coverage_tool = toolchain_attr.configure_coverage_tool,
-        ignore_root_user_error = toolchain_attr.ignore_root_user_error,
-        set_python_version_constraint = version_constraint,
-    )
-    return struct(
-        python_version = toolchain_attr.python_version,
-        set_python_version_constraint = str(version_constraint),
-        name = name,
-    )
-
-def _python_impl(module_ctx):
-    # The toolchain info structs to register, in the order to register them in.
-    toolchains = []
-
-    # We store the default toolchain separately to ensure it is the last
-    # toolchain added to toolchains.
-    default_toolchain = None
-
-    # Map of string Major.Minor to the toolchain name and module name
-    global_toolchain_versions = {}
-
-    for mod in module_ctx.modules:
-        module_toolchain_versions = []
-
-        for toolchain_attr in mod.tags.toolchain:
-            toolchain_version = toolchain_attr.python_version
-            toolchain_name = "python_" + toolchain_version.replace(".", "_")
-
-            # Duplicate versions within a module indicate a misconfigured module.
-            if toolchain_version in module_toolchain_versions:
-                _fail_duplicate_module_toolchain_version(toolchain_version, mod.name)
-            module_toolchain_versions.append(toolchain_version)
-
-            # Ignore version collisions in the global scope because there isn't
-            # much else that can be done. Modules don't know and can't control
-            # what other modules do, so the first in the dependency graph wins.
-            if toolchain_version in global_toolchain_versions:
-                # If the python version is explicitly provided by the root
-                # module, they should not be warned for choosing the same
-                # version that rules_python provides as default.
-                first = global_toolchain_versions[toolchain_version]
-                if mod.name != "rules_python" or not first.is_root:
-                    _warn_duplicate_global_toolchain_version(
-                        toolchain_version,
-                        first = first,
-                        second_toolchain_name = toolchain_name,
-                        second_module_name = mod.name,
-                    )
-                continue
-            global_toolchain_versions[toolchain_version] = struct(
-                toolchain_name = toolchain_name,
-                module_name = mod.name,
-                is_root = mod.is_root,
-            )
-
-            # Only the root module and rules_python are allowed to specify the default
-            # toolchain for a couple reasons:
-            # * It prevents submodules from specifying different defaults and only
-            #   one of them winning.
-            # * rules_python needs to set a soft default in case the root module doesn't,
-            #   e.g. if the root module doesn't use Python itself.
-            # * The root module is allowed to override the rules_python default.
-            if mod.is_root:
-                # A single toolchain is treated as the default because it's unambiguous.
-                is_default = toolchain_attr.is_default or len(mod.tags.toolchain) == 1
-            elif mod.name == "rules_python" and not default_toolchain:
-                # We don't do the len() check because we want the default that rules_python
-                # sets to be clearly visible.
-                is_default = toolchain_attr.is_default
-            else:
-                is_default = False
-
-            # We have already found one default toolchain, and we can only have
-            # one.
-            if is_default and default_toolchain != None:
-                _fail_multiple_default_toolchains(
-                    first = default_toolchain.name,
-                    second = toolchain_name,
-                )
-
-            toolchain_info = _python_register_toolchains(
-                toolchain_name,
-                toolchain_attr,
-                version_constraint = not is_default,
-            )
-
-            if is_default:
-                default_toolchain = toolchain_info
-            else:
-                toolchains.append(toolchain_info)
-
-    # A default toolchain is required so that the non-version-specific rules
-    # are able to match a toolchain.
-    if default_toolchain == None:
-        fail("No default Python toolchain configured. Is rules_python missing `is_default=True`?")
-
-    # The last toolchain in the BUILD file is set as the default
-    # toolchain. We need the default last.
-    toolchains.append(default_toolchain)
-
-    if len(toolchains) > _MAX_NUM_TOOLCHAINS:
-        fail("more than {} python versions are not supported".format(_MAX_NUM_TOOLCHAINS))
-
-    # Create the pythons_hub repo for the interpreter meta data and the
-    # the various toolchains.
-    hub_repo(
-        name = "pythons_hub",
-        default_python_version = default_toolchain.python_version,
-        toolchain_prefixes = [
-            _toolchain_prefix(index, toolchain.name)
-            for index, toolchain in enumerate(toolchains)
-        ],
-        toolchain_python_versions = [t.python_version for t in toolchains],
-        toolchain_set_python_version_constraints = [t.set_python_version_constraint for t in toolchains],
-        toolchain_user_repository_names = [t.name for t in toolchains],
-    )
-
-    # This is require in order to support multiple version py_test
-    # and py_binary
-    multi_toolchain_aliases(
-        name = "python_versions",
-        python_versions = {
-            version: entry.toolchain_name
-            for version, entry in global_toolchain_versions.items()
-        },
-    )
-
-def _fail_duplicate_module_toolchain_version(version, module):
-    fail(("Duplicate module toolchain version: module '{module}' attempted " +
-          "to use version '{version}' multiple times in itself").format(
-        version = version,
-        module = module,
-    ))
-
-def _warn_duplicate_global_toolchain_version(version, first, second_toolchain_name, second_module_name):
-    _print_warn((
-        "Ignoring toolchain '{second_toolchain}' from module '{second_module}': " +
-        "Toolchain '{first_toolchain}' from module '{first_module}' " +
-        "already registered Python version {version} and has precedence"
-    ).format(
-        first_toolchain = first.toolchain_name,
-        first_module = first.module_name,
-        second_module = second_module_name,
-        second_toolchain = second_toolchain_name,
-        version = version,
-    ))
-
-def _fail_multiple_default_toolchains(first, second):
-    fail(("Multiple default toolchains: only one toolchain " +
-          "can have is_default=True. First default " +
-          "was toolchain '{first}'. Second was '{second}'").format(
-        first = first,
-        second = second,
-    ))
-
-python = module_extension(
-    doc = """Bzlmod extension that is used to register Python toolchains.
-""",
-    implementation = _python_impl,
-    tag_classes = {
-        "toolchain": tag_class(
-            doc = """Tag class used to register Python toolchains.
-Use this tag class to register one or more Python toolchains. This class
-is also potentially called by sub modules. The following covers different
-business rules and use cases.
-
-Toolchains in the Root Module
-
-This class registers all toolchains in the root module.
-
-Toolchains in Sub Modules
-
-It will create a toolchain that is in a sub module, if the toolchain
-of the same name does not exist in the root module.  The extension stops name
-clashing between toolchains in the root module and toolchains in sub modules.
-You cannot configure more than one toolchain as the default toolchain.
-
-Toolchain set as the default version
-
-This extension will not create a toolchain that exists in a sub module,
-if the sub module toolchain is marked as the default version. If you have
-more than one toolchain in your root module, you need to set one of the
-toolchains as the default version.  If there is only one toolchain it
-is set as the default toolchain.
-
-Toolchain repository name
-
-A toolchain's repository name uses the format `python_{major}_{minor}`, e.g.
-`python_3_10`. The `major` and `minor` components are
-`major` and `minor` are the Python version from the `python_version` attribute.
-""",
-            attrs = {
-                "configure_coverage_tool": attr.bool(
-                    mandatory = False,
-                    doc = "Whether or not to configure the default coverage tool for the toolchains.",
-                ),
-                "ignore_root_user_error": attr.bool(
-                    default = False,
-                    doc = "Whether the check for root should be ignored or not. This causes cache misses with .pyc files.",
-                    mandatory = False,
-                ),
-                "is_default": attr.bool(
-                    mandatory = False,
-                    doc = "Whether the toolchain is the default version",
-                ),
-                "python_version": attr.string(
-                    mandatory = True,
-                    doc = "The Python version, in `major.minor` format, e.g " +
-                          "'3.12', to create a toolchain for. Patch level " +
-                          "granularity (e.g. '3.12.1') is not supported.",
-                ),
-            },
-        ),
-    },
-)
+python = _python

python/pip_install/BUILD.bazel

@@ -38,6 +38,7 @@ bzl_library(
         "//python/private:render_pkg_aliases_bzl",
         "//python/private:toolchains_repo_bzl",
         "//python/private:which_bzl",
+        "//python/private/bzlmod:pip_repository_bzl",
     ],
 )
 

python/pip_install/pip_repository.bzl

@@ -25,13 +25,17 @@ load("//python/private:normalize_name.bzl", "normalize_name")
 load("//python/private:render_pkg_aliases.bzl", "render_pkg_aliases")
 load("//python/private:toolchains_repo.bzl", "get_host_os_arch")
 load("//python/private:which.bzl", "which_with_fail")
+load("//python/private/bzlmod:pip_repository.bzl", _pip_hub_repository_bzlmod = "pip_repository")
 
 CPPFLAGS = "CPPFLAGS"
 
 COMMAND_LINE_TOOLS_PATH_SLUG = "commandlinetools"
 
 _WHEEL_ENTRY_POINT_PREFIX = "rules_python_wheel_entry_point"
 
+# Kept for not creating merge conflicts with PR#1476, can be removed later.
+pip_hub_repository_bzlmod = _pip_hub_repository_bzlmod
+
 def _construct_pypath(rctx):
     """Helper function to construct a PYTHONPATH.
 
@@ -267,68 +271,6 @@ A requirements_lock attribute must be specified, or a platform-specific lockfile
 """)
     return requirements_txt
 
-def _pip_hub_repository_bzlmod_impl(rctx):
-    bzl_packages = rctx.attr.whl_map.keys()
-    aliases = render_pkg_aliases(
-        repo_name = rctx.attr.repo_name,
-        rules_python = rctx.attr._template.workspace_name,
-        default_version = rctx.attr.default_version,
-        whl_map = rctx.attr.whl_map,
-    )
-    for path, contents in aliases.items():
-        rctx.file(path, contents)
-
-    # NOTE: we are using the canonical name with the double '@' in order to
-    # always uniquely identify a repository, as the labels are being passed as
-    # a string and the resolution of the label happens at the call-site of the
-    # `requirement`, et al. macros.
-    macro_tmpl = "@@{name}//{{}}:{{}}".format(name = rctx.attr.name)
-
-    rctx.file("BUILD.bazel", _BUILD_FILE_CONTENTS)
-    rctx.template("requirements.bzl", rctx.attr._template, substitutions = {
-        "%%ALL_DATA_REQUIREMENTS%%": _format_repr_list([
-            macro_tmpl.format(p, "data")
-            for p in bzl_packages
-        ]),
-        "%%ALL_REQUIREMENTS%%": _format_repr_list([
-            macro_tmpl.format(p, p)
-            for p in bzl_packages
-        ]),
-        "%%ALL_WHL_REQUIREMENTS%%": _format_repr_list([
-            macro_tmpl.format(p, "whl")
-            for p in bzl_packages
-        ]),
-        "%%MACRO_TMPL%%": macro_tmpl,
-        "%%NAME%%": rctx.attr.name,
-    })
-
-pip_hub_repository_bzlmod_attrs = {
-    "default_version": attr.string(
-        mandatory = True,
-        doc = """\
-This is the default python version in the format of X.Y.Z. This should match
-what is setup by the 'python' extension using the 'is_default = True'
-setting.""",
-    ),
-    "repo_name": attr.string(
-        mandatory = True,
-        doc = "The apparent name of the repo. This is needed because in bzlmod, the name attribute becomes the canonical name.",
-    ),
-    "whl_map": attr.string_list_dict(
-        mandatory = True,
-        doc = "The wheel map where values are python versions",
-    ),
-    "_template": attr.label(
-        default = ":pip_repository_requirements_bzlmod.bzl.tmpl",
-    ),
-}
-
-pip_hub_repository_bzlmod = repository_rule(
-    attrs = pip_hub_repository_bzlmod_attrs,
-    doc = """A rule for bzlmod mulitple pip repository creation. PRIVATE USE ONLY.""",
-    implementation = _pip_hub_repository_bzlmod_impl,
-)
-
 def _pip_repository_impl(rctx):
     requirements_txt = locked_requirements_label(rctx, rctx.attr)
     content = rctx.read(requirements_txt)

python/private/BUILD.bazel

@@ -27,6 +27,7 @@ licenses(["notice"])
 filegroup(
     name = "distribution",
     srcs = glob(["**"]) + [
+        "//python/private/bzlmod:distribution",
         "//python/private/common:distribution",
         "//python/private/proto:distribution",
         "//tools/build_defs/python/private:distribution",

python/extensions/private/BUILD.bazel → python/private/bzlmod/BUILD.bazel

@@ -12,12 +12,24 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+load("@bazel_skylib//:bzl_library.bzl", "bzl_library")
+
 package(default_visibility = ["//visibility:private"])
 
 licenses(["notice"])
 
 filegroup(
     name = "distribution",
     srcs = glob(["**"]),
-    visibility = ["//python/extensions/private:__pkg__"],
+    visibility = ["//python/private:__pkg__"],
+)
+
+bzl_library(
+    name = "pip_repository_bzl",
+    srcs = ["pip_repository.bzl"],
+    visibility = ["//:__subpackages__"],
+    deps = [
+        "//python/private:render_pkg_aliases_bzl",
+        "//python/private:text_util_bzl",
+    ],
 )