Platforms/WASI/_build.py

__lazy_modules__ = [
    "contextlib",
    "functools",
    "os",
    "pathlib",
    "shutil",
    "subprocess",
    "sys",
    "tempfile",
]

import contextlib
import functools
import os
import pathlib
import shutil
import subprocess
import sys
import tempfile

try:
    from os import process_cpu_count as cpu_count
except ImportError:
    from os import cpu_count

import _shared

LOCAL_SETUP_MARKER = (
    b"# Generated by Platforms/WASI .\n"
    b"# Required to statically build extension modules."
)

WASMTIME_VAR_NAME = "WASMTIME"
WASMTIME_HOST_RUNNER_VAR = f"{{{WASMTIME_VAR_NAME}}}"


def separator():
    """Print a separator line across the terminal width."""
    try:
        tput_output = subprocess.check_output(
            ["tput", "cols"], encoding="utf-8"
        )
    except subprocess.CalledProcessError:
        terminal_width = 80
    else:
        terminal_width = int(tput_output.strip())
    print("⎯" * terminal_width)


def updated_env(updates={}):
    """Create a new dict representing the environment to use.

    The changes made to the execution environment are printed out.
    """
    env_defaults = {}
    # https://reproducible-builds.org/docs/source-date-epoch/
    git_epoch_cmd = ["git", "log", "-1", "--pretty=%ct"]
    try:
        epoch = subprocess.check_output(
            git_epoch_cmd, encoding="utf-8"
        ).strip()
        env_defaults["SOURCE_DATE_EPOCH"] = epoch
    except subprocess.CalledProcessError:
        pass  # Might be building from a tarball.
    # This layering lets SOURCE_DATE_EPOCH from os.environ takes precedence.
    environment = env_defaults | os.environ | updates

    env_diff = {}
    for key, value in environment.items():
        if os.environ.get(key) != value:
            env_diff[key] = value

    env_vars = [
        f"\n     {key}={item}" for key, item in sorted(env_diff.items())
    ]
    _shared.log("🌎", f"Environment changes:{''.join(env_vars)}")

    return environment


def subdir(context_attr, *, clean_ok=False):
    """Decorator to change to a working directory."""

    def decorator(func):
        @functools.wraps(func)
        def wrapper(context):
            nonlocal context_attr

            working_dir = getattr(context, context_attr)
            separator()
            _shared.log("📁", os.fsdecode(working_dir))
            if clean_ok and context.clean and working_dir.exists():
                _shared.log("🚮", "Deleting directory (--clean)...")
                shutil.rmtree(working_dir)

            working_dir.mkdir(parents=True, exist_ok=True)

            with contextlib.chdir(working_dir):
                return func(context, working_dir)

        return wrapper

    return decorator


def call(command, *, context=None, quiet=False, **kwargs):
    """Execute a command.

    If 'quiet' is true, then redirect stdout and stderr to a temporary file.
    """
    if context is not None:
        quiet = context.quiet

    _shared.log("❯", " ".join(map(str, command)), spacing="  ")
    if not quiet:
        stdout = None
        stderr = None
    else:
        if (log_path := getattr(context, "log_path", None)) is None:
            log_path = pathlib.Path(tempfile.gettempdir())
        stdout = tempfile.NamedTemporaryFile(
            "w",
            encoding="utf-8",
            delete=False,
            dir=log_path,
            prefix="cpython-wasi-",
            suffix=".log",
        )
        stderr = subprocess.STDOUT
        _shared.log("📝", f"Logging output to {stdout.name} (--quiet)...")

    try:
        subprocess.check_call(command, **kwargs, stdout=stdout, stderr=stderr)
    except subprocess.CalledProcessError as error:
        if quiet:
            _shared.log("❌", f"Exit code {error.returncode}")
            separator()
            with open(stdout.name, encoding="utf-8") as file:
                lines = file.readlines()
                # Inefficient, but the log shouldn't be dramatically large.
                print("".join(lines[-10:]), end="")
                if not lines[-1].endswith("\n"):
                    print()
        sys.exit(error.returncode)


@subdir("build_python_path", clean_ok=True)
def configure_build_python(context, working_dir):
    """Configure the build/host Python."""
    if context.setup_local_path.exists():
        if context.setup_local_path.read_bytes() == LOCAL_SETUP_MARKER:
            _shared.log("👍", f"{context.setup_local_path} exists ...")
        else:
            _shared.log(
                "⚠️",
                f"{context.setup_local_path} exists, but has unexpected contents",
            )
    else:
        _shared.log("📝", f"Creating {context.setup_local_path} ...")
        context.setup_local_path.write_bytes(LOCAL_SETUP_MARKER)

    configure = [os.path.relpath(context.checkout / "configure", working_dir)]
    if context.args:
        configure.extend(context.args)

    call(configure, context=context)


@subdir("build_python_path")
def make_build_python(context, _working_dir):
    """Make/build the build Python."""
    call(["make", "--jobs", str(cpu_count()), "all"], context=context)

    binary = context.build_python_interpreter
    cmd = [
        binary,
        "-c",
        "import sys; print(f'{sys.version_info.major}.{sys.version_info.minor}')",
    ]
    version = subprocess.check_output(cmd, encoding="utf-8").strip()

    _shared.log("🎉", f"{binary} {version}")


def wasi_sdk_env(context):
    """Calculate environment variables for building with wasi-sdk."""
    wasi_sdk_path = context.wasi_sdk_path
    sysroot = wasi_sdk_path / "share" / "wasi-sysroot"
    env = {
        "CC": "clang",
        "CPP": "clang-cpp",
        "CXX": "clang++",
        "AR": "llvm-ar",
        "RANLIB": "ranlib",
    }

    for env_var, binary_name in list(env.items()):
        env[env_var] = os.fsdecode(wasi_sdk_path / "bin" / binary_name)

    if not wasi_sdk_path.name.startswith("wasi-sdk"):
        for compiler in ["CC", "CPP", "CXX"]:
            env[compiler] += f" --sysroot={sysroot}"

    env["PKG_CONFIG_PATH"] = ""
    env["PKG_CONFIG_LIBDIR"] = os.pathsep.join(
        map(
            os.fsdecode,
            [sysroot / "lib" / "pkgconfig", sysroot / "share" / "pkgconfig"],
        )
    )
    env["PKG_CONFIG_SYSROOT_DIR"] = os.fsdecode(sysroot)

    env["WASI_SDK_PATH"] = os.fsdecode(wasi_sdk_path)
    env["WASI_SYSROOT"] = os.fsdecode(sysroot)

    env["PATH"] = os.pathsep.join([
        os.fsdecode(wasi_sdk_path / "bin"),
        os.environ["PATH"],
    ])

    return env


@subdir("wasi_build_path", clean_ok=True)
def configure_wasi_python(context, working_dir):
    """Configure the WASI/host build."""
    config_site = os.fsdecode(context.here / "config.site-wasm32-wasi")

    wasi_build_dir = working_dir.relative_to(context.checkout)

    args = {
        "WASMTIME": "wasmtime",
        "ARGV0": f"/{wasi_build_dir}/python.wasm",
        "CHECKOUT": os.fsdecode(context.checkout),
        "WASMTIME_CONFIG_PATH": os.fsdecode(context.here / "wasmtime.toml"),
    }
    # Check dynamically for wasmtime in case it was specified manually via
    # `--host-runner`.
    if "{WASMTIME}" in context.host_runner:
        if wasmtime := shutil.which("wasmtime"):
            args["WASMTIME"] = wasmtime
        else:
            raise FileNotFoundError(
                "wasmtime not found; download from "
                "https://github.com/bytecodealliance/wasmtime"
            )
    host_runner = context.host_runner.format_map(args)
    env_additions = {"CONFIG_SITE": config_site, "HOSTRUNNER": host_runner}
    build_python = os.fsdecode(context.build_python_interpreter)
    # The path to `configure` MUST be relative, else `python.wasm` is unable
    # to find the stdlib due to Python not recognizing that it's being
    # executed from within context.checkout.
    configure = [
        os.path.relpath(context.checkout / "configure", working_dir),
        f"--host={context.host_triple}",
        f"--build={context.build_python_path.name}",
        f"--with-build-python={build_python}",
    ]
    if context.is_debug:
        configure.append("--with-pydebug")
    if context.args:
        configure.extend(context.args)
    call(
        configure,
        env=updated_env(env_additions | wasi_sdk_env(context)),
        context=context,
    )

    python_wasm = working_dir / "python.wasm"
    exec_script = working_dir / "python.sh"
    with exec_script.open("w", encoding="utf-8") as file:
        file.write(f'#!/bin/sh\nexec {host_runner} {python_wasm} "$@"\n')
    exec_script.chmod(0o755)
    _shared.log("🏃", f"Created {exec_script} (--host-runner)... ")
    sys.stdout.flush()


@subdir("wasi_build_path")
def make_wasi_python(context, working_dir):
    """Run `make` for the WASI/host build."""
    call(
        ["make", "--jobs", str(cpu_count()), "all"],
        env=updated_env(),
        context=context,
    )

    exec_script = working_dir / "python.sh"
    call([exec_script, "-c", "import sys; print(sys.version)"], quiet=False)
    _shared.log(
        "🎉",
        f"Use `{exec_script.relative_to(pathlib.Path().absolute())}` "
        "to run CPython w/ the WASI host specified by --host-runner",
    )


def clean_contents(context):
    """Delete all files created by this script."""
    context.clean = True
    if context.cross_build_path.exists():
        _shared.log("🧹", f"Deleting {context.cross_build_path} ...")
        shutil.rmtree(context.cross_build_path)

    if context.setup_local_path.exists():
        if context.setup_local_path.read_bytes() == LOCAL_SETUP_MARKER:
            _shared.log(
                "🧹", f"Deleting generated {context.setup_local_path} ..."
            )
            context.setup_local_path.unlink()


@subdir("build_python_path")
def pythoninfo_build_python(context, working_dir):
    """Display build info of the build Python."""
    call(["make", "pythoninfo"], context=context)


@subdir("wasi_build_path")
def pythoninfo_wasi_python(context, working_dir):
    """Display build info of the host/WASI Python."""
    call(["make", "pythoninfo"], context=context)