local_runtime_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 for a locally installed Python runtime."""

load(":enum.bzl", "enum")
load(":repo_utils.bzl", "REPO_DEBUG_ENV_VAR", "repo_utils")

# buildifier: disable=name-conventions
_OnFailure = enum(
    SKIP = "skip",
    WARN = "warn",
    FAIL = "fail",
)

_TOOLCHAIN_IMPL_TEMPLATE = """\
# Generated by python/private/local_runtime_repo.bzl

load("@rules_python//python/private:local_runtime_repo_setup.bzl", "define_local_runtime_toolchain_impl")

define_local_runtime_toolchain_impl(
    name = "local_runtime",
    major = "{major}",
    minor = "{minor}",
    micro = "{micro}",
    abi_flags = "{abi_flags}",
    os = "{os}",
    implementation_name = "{implementation_name}",
    interpreter_path = "{interpreter_path}",
    interface_library = {interface_library},
    libraries = {libraries},
    defines = {defines},
    abi3_interface_library = {abi3_interface_library},
    abi3_libraries = {abi3_libraries},
    additional_dlls = {additional_dlls},
)
"""

def _expand_incompatible_template():
    return _TOOLCHAIN_IMPL_TEMPLATE.format(
        major = "0",
        minor = "0",
        micro = "0",
        abi_flags = "",
        os = "@platforms//:incompatible",
        implementation_name = "incompatible",
        interpreter_path = "/incompatible",
        interface_library = "None",
        libraries = "[]",
        defines = "[]",
        abi3_interface_library = "None",
        abi3_libraries = "[]",
        additional_dlls = "[]",
    )

def _norm_path(path):
    """Returns a path using '/' separators and no trailing slash."""
    path = path.replace("\\", "/")
    if path[-1] == "/":
        path = path[:-1]
    return path

def _symlink_libraries(rctx, logger, libraries, shlib_suffix):
    """Symlinks the shared libraries into the lib/ directory.

    Individual files are symlinked instead of the whole directory because
    shared_lib_dirs contains multiple search paths for the shared libraries,
    and the python files may be missing from any of those directories, and
    any of those directories may include non-python runtime libraries,
    as would be the case if LIBDIR were, for example, /usr/lib.

    Args:
        rctx: A repository_ctx object
        logger: A repo_utils.logger object
        libraries: paths to libraries to attempt to symlink.
        shlib_suffix: Optional. Ensure that the generated symlinks end with this suffix.
    Returns:
        A list of library paths (under lib/) linked by the action.
    """
    result = []
    for source in libraries:
        origin = rctx.path(source)
        if not origin.exists:
            # The reported names don't always exist; it depends on the particulars
            # of the runtime installation.
            continue
        if shlib_suffix and not origin.basename.endswith(shlib_suffix):
            target = "lib/{}{}".format(origin.basename, shlib_suffix)
        else:
            target = "lib/{}".format(origin.basename)
        logger.debug(lambda: "Symlinking {} to {}".format(origin, target))
        rctx.watch(origin)
        rctx.symlink(origin, target)
        result.append(target)
    return result

def _local_runtime_repo_impl(rctx):
    logger = repo_utils.logger(rctx)
    on_failure = rctx.attr.on_failure

    def _emit_log(msg):
        if on_failure == "fail":
            logger.fail(msg)
        elif on_failure == "warn":
            logger.warn(msg)
        else:
            logger.debug(msg)

    result = _resolve_interpreter_path(rctx)
    if not result.resolved_path:
        _emit_log(lambda: "interpreter not found: {}".format(result.describe_failure()))

        # else, on_failure must be skip
        rctx.file("BUILD.bazel", _expand_incompatible_template())
        return
    else:
        interpreter_path = result.resolved_path

    logger.info(lambda: "resolved interpreter {} to {}".format(rctx.attr.interpreter_path, interpreter_path))

    exec_result = repo_utils.execute_unchecked(
        rctx,
        op = "local_runtime_repo.GetPythonInfo({})".format(rctx.name),
        arguments = [
            interpreter_path,
            rctx.path(rctx.attr._get_local_runtime_info),
        ],
        quiet = True,
        logger = logger,
    )
    if exec_result.return_code != 0:
        _emit_log(lambda: "GetPythonInfo failed: {}".format(exec_result.describe_failure()))

        # else, on_failure must be skip
        rctx.file("BUILD.bazel", _expand_incompatible_template())
        return

    info = json.decode(exec_result.stdout)
    logger.info(lambda: _format_get_info_result(info))

    # We use base_executable because we want the path within a Python
    # installation directory ("PYTHONHOME"). The problems with sys.executable
    # are:
    # * If we're in an activated venv, then we don't want the venv's
    #   `bin/python3` path to be used -- it isn't an actual Python installation.
    # * If sys.executable is a wrapper (e.g. pyenv), then (1) it may not be
    #   located within an actual Python installation directory, and (2) it
    #   can interfer with Python recognizing when it's within a venv.
    #
    # In some cases, it may be a symlink (usually e.g. `python3->python3.12`),
    # but we don't realpath() it to respect what it has decided is the
    # appropriate path.
    interpreter_path = info["base_executable"]

    # NOTE: Keep in sync with recursive glob in define_local_runtime_toolchain_impl
    include_path = rctx.path(info["include"])

    # The reported include path may not exist, and watching a non-existant
    # path is an error. Silently skip, since includes are only necessary
    # if C extensions are built.
    if include_path.exists and include_path.is_dir:
        rctx.watch_tree(include_path)
    else:
        pass

    # The cc_library.includes values have to be non-absolute paths, otherwise
    # the toolchain will give an error. Work around this error by making them
    # appear as part of this repo.
    logger.debug(lambda: "Symlinking {} to include".format(include_path))
    rctx.symlink(include_path, "include")

    rctx.report_progress("Symlinking external Python shared libraries")

    interface_library = None
    if info["dynamic_libraries"]:
        libraries = _symlink_libraries(rctx, logger, info["dynamic_libraries"][:1], info["shlib_suffix"])
        symlinked = _symlink_libraries(rctx, logger, info["interface_libraries"][:1], None)
        if symlinked:
            interface_library = symlinked[0]
    else:
        libraries = _symlink_libraries(rctx, logger, info["static_libraries"], None)
        if not libraries:
            logger.info("No python libraries found.")

    abi3_interface_library = None
    if info["abi_dynamic_libraries"]:
        abi3_libraries = _symlink_libraries(rctx, logger, info["abi_dynamic_libraries"][:1], info["shlib_suffix"])
        symlinked = _symlink_libraries(rctx, logger, info["abi_interface_libraries"][:1], None)
        if symlinked:
            abi3_interface_library = symlinked[0]
    else:
        abi3_libraries = []
        logger.info("No abi3 python libraries found.")

    additional_dlls = _symlink_libraries(rctx, logger, info["additional_dlls"], None)

    build_bazel = _TOOLCHAIN_IMPL_TEMPLATE.format(
        major = info["major"],
        minor = info["minor"],
        micro = info["micro"],
        abi_flags = info["abi_flags"],
        os = "@platforms//os:{}".format(repo_utils.get_platforms_os_name(rctx)),
        implementation_name = info["implementation_name"],
        interpreter_path = _norm_path(interpreter_path),
        interface_library = repr(interface_library),
        libraries = repr(libraries),
        defines = repr(info["defines"]),
        abi3_interface_library = repr(abi3_interface_library),
        abi3_libraries = repr(abi3_libraries),
        additional_dlls = repr(additional_dlls),
    )
    logger.debug(lambda: "BUILD.bazel\n{}".format(build_bazel))

    rctx.file("WORKSPACE", "")
    rctx.file("MODULE.bazel", "")
    rctx.file("REPO.bazel", "")
    rctx.file("BUILD.bazel", build_bazel)

local_runtime_repo = repository_rule(
    implementation = _local_runtime_repo_impl,
    doc = """
Use a locally installed Python runtime as a toolchain implementation.

Note this uses the runtime as a *platform runtime*. A platform runtime means
means targets don't include the runtime itself as part of their runfiles or
inputs. Instead, users must assure that where the targets run have the runtime
pre-installed or otherwise available.

This results in lighter weight binaries (in particular, Bazel doesn't have to
create thousands of files for every `py_test`), at the risk of having to rely on
a system having the necessary Python installed.
""",
    attrs = {
        "interpreter_path": attr.string(
            doc = """
An absolute path or program name on the `PATH` env var.

*Mutually exclusive with `interpreter_target`.*

Values with slashes are assumed to be the path to a program. Otherwise, it is
treated as something to search for on `PATH`

Note that, when a plain program name is used, the path to the interpreter is
resolved at repository evalution time, not runtime of any resulting binaries.

If not set, defaults to `python3`.

:::{seealso}
The {obj}`interpreter_target` attribute for getting the interpreter from
a label
:::
""",
            default = "",
        ),
        "interpreter_target": attr.label(
            doc = """
A label to a Python interpreter executable.

*Mutually exclusive with `interpreter_path`.*

On Windows, if the path doesn't exist, various suffixes will be tried to
find a usable path.

:::{seealso}
The {obj}`interpreter_path` attribute for getting the interpreter from
a path or PATH environment lookup.
:::
""",
        ),
        "on_failure": attr.string(
            default = _OnFailure.SKIP,
            values = sorted(_OnFailure.__members__.values()),
            doc = """
How to handle errors when trying to automatically determine settings.

* `skip` will silently skip creating a runtime. Instead, a non-functional
  runtime will be generated and marked as incompatible so it cannot be used.
  This is best if a local runtime is known not to work or be available
  in certain cases and that's OK. e.g., one use windows paths when there
  are people running on linux.
* `warn` will print a warning message. This is useful when you expect
  a runtime to be available, but are OK with it missing and falling back
  to some other runtime.
* `fail` will result in a failure. This is only recommended if you must
  ensure the runtime is available.
""",
        ),
        "_get_local_runtime_info": attr.label(
            allow_single_file = True,
            default = "//python/private:get_local_runtime_info.py",
        ),
        "_rule_name": attr.string(default = "local_runtime_repo"),
    },
    environ = ["PATH", REPO_DEBUG_ENV_VAR, "DEVELOPER_DIR", "XCODE_VERSION"],
)

def _find_python_exe_from_target(rctx):
    base_path = rctx.path(rctx.attr.interpreter_target)
    if base_path.exists:
        return base_path, None
    attempted_paths = [base_path]

    # Try to convert a unix-y path to a Windows path. On Linux/Mac,
    # the path is usually `bin/python3`. On Windows, it's simply
    # `python.exe`.
    basename = base_path.basename.rstrip("3")
    path = base_path.dirname.dirname.get_child(basename)
    path = rctx.path("{}.exe".format(path))
    if path.exists:
        return path, None
    attempted_paths.append(path)

    # Try adding .exe to the base path
    path = rctx.path("{}.exe".format(base_path))
    if path.exists:
        return path, None
    attempted_paths.append(path)

    describe_failure = lambda: (
        "Target '{target}' could not be resolved to a valid path. " +
        "Attempted paths: {paths}"
    ).format(
        target = rctx.attr.interpreter_target,
        paths = "\n".join([str(p) for p in attempted_paths]),
    )
    return None, describe_failure

def _resolve_interpreter_path(rctx):
    """Find the absolute path for an interpreter.

    Args:
        rctx: A repository_ctx object

    Returns:
        `struct` with the following fields:
        * `resolved_path`: `path` object of a path that exists
        * `describe_failure`: `Callable | None`. If a path that doesn't exist,
          returns a description of why it couldn't be resolved
        A path object or None. The path may not exist.
    """
    if rctx.attr.interpreter_path and rctx.attr.interpreter_target:
        fail("interpreter_path and interpreter_target are mutually exclusive")

    if rctx.attr.interpreter_target:
        resolved_path, describe_failure = _find_python_exe_from_target(rctx)
    else:
        interpreter_path = rctx.attr.interpreter_path or "python3"
        if "/" not in interpreter_path and "\\" not in interpreter_path:
            # Provide a bit nicer integration with pyenv: recalculate the runtime if the
            # user changes the python version using e.g. `pyenv shell`
            rctx.getenv("PYENV_VERSION")
            result = repo_utils.which_unchecked(rctx, interpreter_path)
            resolved_path = result.binary
            describe_failure = result.describe_failure
        else:
            rctx.watch(interpreter_path)
            resolved_path = rctx.path(interpreter_path)
            if not resolved_path.exists:
                describe_failure = lambda: "Path not found: {}".format(repr(interpreter_path))
            else:
                describe_failure = None

    return struct(
        resolved_path = resolved_path,
        describe_failure = describe_failure,
    )

def _format_get_info_result(info):
    lines = ["GetPythonInfo result:"]
    for key, value in sorted(info.items()):
        lines.append("  {}: {}".format(key, value if value != "" else "<empty string>"))
    return "\n".join(lines)