local_runtime_toolchains_repo.bzl

# 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.

"""Create a repository to hold a local Python toolchain definitions."""

load(":repo_utils.bzl", "REPO_DEBUG_ENV_VAR", "repo_utils")
load(":text_util.bzl", "render")

_TOOLCHAIN_TEMPLATE = """
# Generated by local_runtime_toolchains_repo.bzl

load("@rules_python//python/private:py_toolchain_suite.bzl", "define_local_toolchain_suites")

define_local_toolchain_suites(
    name = "toolchains",
    version_aware_repo_names = {version_aware_names},
    version_unaware_repo_names = {version_unaware_names},
)
"""

def _local_runtime_toolchains_repo(rctx):
    logger = repo_utils.logger(rctx)
    rctx.file("WORKSPACE", "")
    rctx.file("MODULE.bazel", "")
    rctx.file("REPO.bazel", "")

    logger.info(lambda: _format_toolchains_for_logging(rctx))

    rctx.file("BUILD.bazel", _TOOLCHAIN_TEMPLATE.format(
        version_aware_names = render.list(rctx.attr.runtimes),
        version_unaware_names = render.list(rctx.attr.default_runtimes or rctx.attr.runtimes),
    ))

local_runtime_toolchains_repo = repository_rule(
    implementation = _local_runtime_toolchains_repo,
    doc = """
Create a repo of toolchains definitions for local runtimes.

This is intended to be used on the toolchain implemenations generated by
`local_runtime_repo`.

NOTE: This does not call `native.register_toolchains` -- the caller is
responsible for registering the toolchains this defines.
""",
    attrs = {
        "default_runtimes": attr.string_list(
            doc = """
The repo names of `local_runtime_repo` repos to define as toolchains.

These will be defined as *version-unaware* toolchains. This means they will
match any Python version. As such, they are registered after the version-aware
toolchains defined by the `runtimes` attribute.

Note that order matters: it determines the toolchain priority within the
package.
""",
        ),
        "runtimes": attr.string_list(
            doc = """
The repo names of `local_runtime_repo` repos to define as toolchains.

These will be defined as *version-aware* toolchains. This means they require the
`--//python/config_settings:python_version` to be set in order to match. These
are registered before `default_runtimes`.

Note that order matters: it determines the toolchain priority within the
package.
""",
        ),
        "_rule_name": attr.string(default = "local_toolchains_repo"),
    },
    environ = [REPO_DEBUG_ENV_VAR],
)

def _format_toolchains_for_logging(rctx):
    lines = ["Local toolchain priority order:"]
    i = 0
    for i, name in enumerate(rctx.attr.runtimes, start = i):
        lines.append("  {}: {} (version aware)".format(i, name))
    for i, name in enumerate(rctx.attr.default_runtimes, start = i):
        lines.append("  {}: {} (version unaware)".format(i, name))
    return "\n".join(lines)