hatch_build.py

import subprocess
import shutil
import sys
import os
from packaging.tags import sys_tags
from hatchling.builders.hooks.plugin.interface import BuildHookInterface
from pathlib import Path
import json


class CargoWorkspaceBuild:
    def __init__(self, hook: "BundleBuildHook") -> None:
        self.hook = hook
        self.metadata = self._get_metadata()

    def _get_metadata(self):
        p = subprocess.Popen(
            [
                "cargo",
                "metadata",
                "--no-deps",
            ],
            cwd=self.hook.root,
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE,
        )
        stdout, stderr = p.communicate()
        if p.returncode != 0:
            self.hook.app.display_error(
                f"Cargo metadata command failed with return code {p.returncode}"
            )
            self.hook.app.display_error(stderr.decode("utf-8"))
            sys.exit(1)
        return json.loads(stdout)

    def _run_build(self, args: list[str], env: dict[str, str] | None = None):
        p = subprocess.Popen(
            args,
            cwd=self.hook.root,
            stdout=subprocess.PIPE,
            stderr=subprocess.STDOUT,
            env=env,
        )
        assert p.stdout is not None
        for line in p.stdout:
            line_str = line.decode("utf-8").strip()
            if "finished" in line_str.lower():
                self.hook.app.display_info(line_str)
            elif "error" in line_str.lower():
                self.hook.app.display_error(line_str)
            else:
                self.hook.app.display_info(line_str)
        p.wait()
        if p.returncode != 0:
            self.hook.app.display_error(
                f"Cargo build failed with return code {p.returncode}"
            )
            sys.exit(1)

    def build_all(self):
        self.hook.app.display_mini_header("Building cargo workspace")
        self._run_build(
            [
                "cargo",
                "build",
                "--workspace",
                "--release",
                "--locked",
            ]
        )
        self.hook.app.display_success("Cargo build completed successfully")

    def extract_libs(self):
        release_dir = self.hook.get_cargo_release_dir()
        assert release_dir.exists()
        for package in self.metadata["packages"]:
            for target in package["targets"]:
                if "cdylib" not in target["kind"]:
                    continue
                self.hook.app.display_info(f"Found cdylib target: {target['name']}")
                lib_name = target["name"]
                lib_filenames = {
                    "darwin": [f"lib{lib_name}.dylib"],
                    "linux": [f"lib{lib_name}.so"],
                    "win32": [f"{lib_name}.dll", f"lib{lib_name}.dll.a"],
                }[sys.platform]
                assert all((release_dir / file).exists() for file in lib_filenames), (
                    f"Compiled library for {lib_name} not found in {release_dir}. "
                    f"Expected to find {','.join(lib_filenames)}."
                )
                self.hook.app.display_info(f"Found {lib_name} in {release_dir}")

                cargo_toml_path = Path(package["manifest_path"])
                python_path = cargo_toml_path.parent / "python"
                subdirs = filter(lambda x: x.is_dir(), python_path.iterdir())
                subdirs = filter(
                    lambda x: x.name != "test" and x.name != "tests", subdirs
                )
                subdirs = filter(lambda x: "pycache" not in x.name, subdirs)
                subdirs = filter(lambda x: not x.name.endswith("egg-info"), subdirs)
                subdirs = filter(lambda x: not x.name.startswith("."), subdirs)
                subdirs_list = list(subdirs)
                assert len(subdirs_list) == 1, (
                    f"Multiple python directories found in {python_path} - "
                    " hatch_build.py can't tell where the build artifacts should go."
                )
                subdir = subdirs_list[0]
                destination = subdir / "_dist/lib"
                destination.mkdir(parents=True, exist_ok=True)
                for filename in lib_filenames:
                    lib_path = release_dir / filename
                    if lib_path.exists():
                        self.hook.app.display_info(
                            f"Copying {lib_path} to {destination}"
                        )
                        shutil.copy(lib_path, destination)


class UtilitiesBuild:
    def __init__(self, hook: "BundleBuildHook") -> None:
        self.hook = hook
        self.utilities = list(Path(self.hook.root).glob("selene-ext/utilities/*/"))

    def _get_release_dir(self, utility_root: Path) -> Path:
        target = os.environ.get("CARGO_BUILD_TARGET")
        target_dir = utility_root / "target"
        if target:
            return target_dir / target / "release"
        return target_dir / "release"

    def _get_metadata(self, utility_root: Path) -> dict:
        p = subprocess.Popen(
            [
                "cargo",
                "metadata",
                "--no-deps",
                "--manifest-path",
                str(utility_root / "Cargo.toml"),
            ],
            cwd=self.hook.root,
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE,
        )
        stdout, stderr = p.communicate()
        if p.returncode != 0:
            self.hook.app.display_error(
                f"Cargo metadata command failed with return code {p.returncode}"
            )
            self.hook.app.display_error(stderr.decode("utf-8"))
            sys.exit(1)
        return json.loads(stdout)

    def build_all(self):
        if not self.utilities:
            return
        env = os.environ.copy()
        env["SELENE_BASE_QIS_LIB_DIR"] = str(
            Path(self.hook.root)
            / "selene-ext/interfaces/base_qis/python/selene_base_qis_plugin/_dist/lib"
        )
        env["SELENE_HELIOS_QIS_LIB_DIR"] = str(
            Path(self.hook.root)
            / "selene-ext/interfaces/helios_qis/python/selene_helios_qis_plugin/_dist/lib"
        )
        env["SELENE_SOL_QIS_LIB_DIR"] = str(
            Path(self.hook.root)
            / "selene-ext/interfaces/sol_qis/python/selene_sol_qis_plugin/_dist/lib"
        )
        for utility in self.utilities:
            self.hook.app.display_mini_header(f"Building utility {utility}")
            p = subprocess.Popen(
                [
                    "cargo",
                    "build",
                    "--manifest-path",
                    str(utility / "Cargo.toml"),
                    "--release",
                    # "--locked",
                ],
                cwd=self.hook.root,
                stdout=subprocess.PIPE,
                stderr=subprocess.STDOUT,
                env=env,
            )
            assert p.stdout is not None
            for line in p.stdout:
                line_str = line.decode("utf-8").strip()
                if "finished" in line_str.lower():
                    self.hook.app.display_info(line_str)
                elif "error" in line_str.lower():
                    self.hook.app.display_error(line_str)
                else:
                    self.hook.app.display_info(line_str)
            p.wait()
            if p.returncode != 0:
                self.hook.app.display_error(
                    f"Utility build failed with return code {p.returncode}"
                )
                sys.exit(1)
            self.hook.app.display_success(
                f"Utility build completed successfully: {utility}"
            )

    def extract_libs(self):
        for utility in self.utilities:
            metadata = self._get_metadata(utility)
            release_dir = self._get_release_dir(utility)
            assert release_dir.exists()
            for package in metadata["packages"]:
                for target in package["targets"]:
                    if "cdylib" not in target["kind"]:
                        continue
                    self.hook.app.display_info(
                        f"Found addon cdylib target: {target['name']}"
                    )
                    lib_name = target["name"]
                    lib_filenames = {
                        "darwin": [f"lib{lib_name}.dylib"],
                        "linux": [f"lib{lib_name}.so"],
                        "win32": [f"{lib_name}.dll", f"lib{lib_name}.dll.a"],
                    }[sys.platform]
                    assert all(
                        (release_dir / file).exists() for file in lib_filenames
                    ), (
                        f"Compiled library for {lib_name} not found in {release_dir}. "
                        f"Expected to find {','.join(lib_filenames)}."
                    )

                    cargo_toml_path = Path(package["manifest_path"])
                    python_path = cargo_toml_path.parent / "python"
                    subdirs = filter(lambda x: x.is_dir(), python_path.iterdir())
                    subdirs = filter(
                        lambda x: x.name != "test" and x.name != "tests", subdirs
                    )
                    subdirs = filter(lambda x: "pycache" not in x.name, subdirs)
                    subdirs = filter(lambda x: not x.name.endswith("egg-info"), subdirs)
                    subdirs = filter(lambda x: not x.name.startswith("."), subdirs)
                    subdirs_list = list(subdirs)
                    assert len(subdirs_list) == 1, (
                        f"Multiple python directories found in {python_path} - "
                        " hatch_build.py can't tell where the build artifacts should go."
                    )
                    subdir = subdirs_list[0]
                    destination = subdir / "_dist/lib"
                    destination.mkdir(parents=True, exist_ok=True)
                    for filename in lib_filenames:
                        lib_path = release_dir / filename
                        if lib_path.exists():
                            self.hook.app.display_info(
                                f"Copying {lib_path} to {destination}"
                            )
                            shutil.copy(lib_path, destination)


class BundleBuildHook(BuildHookInterface):
    def target_is_windows(self) -> bool:
        return sys.platform == "win32" or os.environ.get(
            "CARGO_BUILD_TARGET", ""
        ).endswith("windows-gnu")

    def target_is_windows_gnu(self) -> bool:
        return os.environ.get("CARGO_BUILD_TARGET", "").endswith("windows-gnu")

    def runtime_origin(self) -> str:
        if sys.platform == "darwin":
            return "@loader_path"
        return "$ORIGIN"

    def install_rpath_arg(self, rpaths: list[str]) -> list[str]:
        if self.target_is_windows():
            return []
        return [f"-DCMAKE_INSTALL_RPATH={';'.join(rpaths)}"]

    def get_cargo_release_dir(self) -> Path:
        target = os.environ.get("CARGO_BUILD_TARGET")
        if target:
            candidate = Path(self.root) / "target" / target / "release"
            if candidate.exists():
                return candidate
        return Path(self.root) / "target/release"

    def initialize(self, version: str, build_data: dict) -> None:
        cargo_runner = CargoWorkspaceBuild(self)
        cargo_runner.build_all()
        cargo_runner.extract_libs()
        self.build_selene_c_interface()
        self.build_base_qis()
        self.build_platform_qis("helios")
        self.build_platform_qis("sol")
        utilities_builder = UtilitiesBuild(self)
        utilities_builder.build_all()
        utilities_builder.extract_libs()

        packages = [Path("selene-sim/python/selene_sim")]
        for topic_dir in Path("selene-ext").iterdir():
            if topic_dir.name.startswith("."):
                continue
            if not topic_dir.is_dir():
                continue
            for package in topic_dir.iterdir():
                if package.name.startswith("."):
                    continue
                if not package.is_dir():
                    continue
                python_dir = package / "python"
                if not python_dir.exists():
                    continue
                subdirs = filter(lambda x: x.is_dir(), python_dir.iterdir())
                subdirs = filter(
                    lambda x: x.name != "test" and x.name != "tests", subdirs
                )
                subdirs = filter(lambda x: "pycache" not in x.name, subdirs)
                subdirs = filter(lambda x: not x.name.endswith("egg-info"), subdirs)
                subdirs = filter(lambda x: not x.name.startswith("."), subdirs)
                subdirs_list = list(subdirs)
                if len(subdirs_list) == 0:
                    continue
                if len(subdirs_list) > 1:
                    self.app.display_error(
                        f"Multiple python directories found in {python_dir} - "
                        "hatch_build.py can't tell which one to use."
                    )
                    sys.exit(1)
                package = subdirs_list[0]
                packages.append(package)
                self.app.display_info(f"Found package: {package}")

        artifacts = []
        for package in packages:
            package_root = Path(package)
            # add py.typed
            py_typed = package_root / "py.typed"
            if py_typed.exists():
                artifacts.append(str(py_typed.as_posix()))
            # add distribution files (e.g. dynamic libraries, headers)
            dist_dir = Path(package) / "_dist"
            if not dist_dir.exists():
                pass
            for artifact in dist_dir.rglob("*"):
                if artifact.is_file():
                    artifacts.append(str(artifact.as_posix()))

        self.app.display_info("Found artifacts:")
        for a in artifacts:
            self.app.display_info(f"    {a}")

        build_data["packages"] = packages
        build_data["artifacts"] += artifacts
        build_data["pure_python"] = False
        # Set platform-specific wheel tags
        # the approach is an alternative of
        # https://github.com/pypa/hatch/blob/9e1fc3472f9f2536e9269cd2009f878e597a6061/backend/src/hatchling/builders/wheel.py#L782
        # but does not use the interpreter or ABI components, as selene's compiled
        # libs do not bind to python itself.
        tag = next(
            iter(
                t
                for t in sys_tags()
                if "manylinux" not in t.platform and "musllinux" not in t.platform
            )
        )
        target_platform = tag.platform
        if sys.platform == "darwin":
            from hatchling.builders.macos import process_macos_plat_tag

            target_platform = process_macos_plat_tag(target_platform, compat=False)
        build_data["tag"] = f"py3-none-{target_platform}"

    def find_release_files(self, cdylib_name):
        release_dir = self.get_cargo_release_dir()
        if sys.platform == "darwin":
            return [release_dir / f"lib{cdylib_name}.dylib"]
        elif sys.platform == "linux":
            return [release_dir / f"lib{cdylib_name}.so"]
        elif sys.platform == "win32":
            return [
                release_dir / f"{cdylib_name}.dll",
                release_dir / f"lib{cdylib_name}.dll.a",
            ]

    def build_cargo_workspace(self):
        self.app.display_mini_header("Building cargo workspace")
        try:
            subprocess.run(
                [
                    "cargo",
                    "build",
                    "--workspace",
                    "--release",
                    "--locked",
                ],
                cwd=self.root,
                check=True,
                capture_output=True,
            )
        except subprocess.CalledProcessError as e:
            self.app.display_error(f"Cargo build failed: {e.stderr.decode()}")
            sys.exit(1)
        self.app.display_success("Cargo build completed successfully")

    def distribute_cargo_artifacts(self):
        self.app.display_waiting("Copying artifacts")
        lib_paths = []
        for lib in self.find_release_files("selene"):
            self.app.display_info(f"Copying {lib} to {self.install_lib_dir}")
            shutil.copy(lib, self.install_lib_dir)
            lib_paths.append(lib)
        for lib in self.find_release_files("helios_selene_interface"):
            self.app.display_info(f"Copying {lib} to {self.install_lib_dir}")
            shutil.copy(lib, self.install_lib_dir)
            lib_paths.append(lib)
        for lib in self.find_release_files("sol_selene_interface"):
            self.app.display_info(f"Copying {lib} to {self.install_lib_dir}")
            shutil.copy(lib, self.install_lib_dir)
            lib_paths.append(lib)

        return lib_paths

    def build_selene_c_interface(self):
        self.app.display_mini_header("Building Selene C interface")
        selene_sim_dir = Path(self.root) / "selene-sim"
        dist_dir = selene_sim_dir / "python/selene_sim/_dist"
        dist_dir.mkdir(parents=True, exist_ok=True)

        cmake_source_dir = selene_sim_dir / "c"

        cmake_build_dir = Path(self.root) / "target" / "selene_c_interface_build"
        cmake_build_dir.mkdir(parents=True, exist_ok=True)

        try:
            subprocess.run(
                [
                    "cmake",
                    f"-DCMAKE_INSTALL_PREFIX={dist_dir}",
                    "-DCMAKE_BUILD_TYPE=Release",
                    f"{cmake_source_dir}",
                ],
                cwd=cmake_build_dir,
                check=True,
                capture_output=True,
            )
        except subprocess.CalledProcessError as e:
            if b"is different than the directory" not in e.stderr:
                self.app.display_error(f"cmake failed: {e.stderr.decode()}")
                sys.exit(1)
            try:
                # existing build dir is incompatible, delete and retry
                shutil.rmtree(cmake_build_dir)
                cmake_build_dir.mkdir()
                subprocess.run(
                    [
                        "cmake",
                        f"-DCMAKE_INSTALL_PREFIX={dist_dir}",
                        "-DCMAKE_BUILD_TYPE=Release",
                        f"{cmake_source_dir}",
                    ],
                    cwd=cmake_build_dir,
                    check=True,
                    capture_output=True,
                )
            except subprocess.CalledProcessError as e:
                self.app.display_error(f"cmake failed: {e.stderr.decode()}")
                sys.exit(1)

        try:
            subprocess.run(
                [
                    "cmake",
                    "--build",
                    ".",
                    "--target",
                    "install",
                ],
                cwd=cmake_build_dir,
                check=True,
                capture_output=True,
            )
        except subprocess.CalledProcessError as e:
            self.app.display_error(f"cmake build failed: {e.stderr.decode()}")
            sys.exit(1)

        self.app.display_success("C interface build completed successfully")

    def build_base_qis(self):
        self.app.display_mini_header("Building base QIS")
        base_qis_dir = Path(self.root) / "selene-ext/interfaces/base_qis"
        cmake_source_dir = base_qis_dir / "c"
        cmake_build_dir = Path(self.root) / "target" / "base_qis_build"
        cmake_build_dir.mkdir(parents=True, exist_ok=True)
        dist_dir = base_qis_dir / "python/selene_base_qis_plugin/_dist"
        dist_dir.mkdir(parents=True, exist_ok=True)
        selene_sim_dist_dir = Path(self.root) / "selene-sim/python/selene_sim/_dist"
        local_dist_lib_dir = dist_dir / "lib"
        local_selene_dist_lib_dir = selene_sim_dist_dir / "lib"
        local_relative_path = os.path.relpath(
            local_selene_dist_lib_dir.resolve(), local_dist_lib_dir.resolve()
        )
        origin = self.runtime_origin()
        local_rpath = (
            origin if local_relative_path == "." else f"{origin}/{local_relative_path}"
        )
        installed_rpath = f"{origin}/../../../selene_sim/_dist/lib"
        cmake_configure_cmd = [
            "cmake",
            "-DCMAKE_INSTALL_LIBDIR=lib",
            f"-DCMAKE_INSTALL_PREFIX={dist_dir}",
            *self.install_rpath_arg([local_rpath, installed_rpath]),
            "-DCMAKE_BUILD_TYPE=Release",
            f"-DCMAKE_PREFIX_PATH={selene_sim_dist_dir}",
            f"{cmake_source_dir}",
        ]
        if self.target_is_windows_gnu():
            cmake_configure_cmd = [
                cmake_configure_cmd[0],
                "-G",
                "MinGW Makefiles",
                *cmake_configure_cmd[1:],
            ]

        try:
            subprocess.run(
                cmake_configure_cmd,
                cwd=cmake_build_dir,
                check=True,
                capture_output=True,
            )
        except subprocess.CalledProcessError as e:
            if b"is different than the directory" not in e.stderr:
                self.app.display_error(f"cmake failed: {e.stderr.decode()}")
                sys.exit(1)
            try:
                # existing build dir is incompatible, delete and retry
                shutil.rmtree(cmake_build_dir)
                cmake_build_dir.mkdir()
                subprocess.run(
                    cmake_configure_cmd,
                    cwd=cmake_build_dir,
                    check=True,
                    capture_output=True,
                )
            except subprocess.CalledProcessError as e:
                self.app.display_error(f"cmake failed: {e.stderr.decode()}")
                sys.exit(1)

        try:
            subprocess.run(
                [
                    "cmake",
                    "--build",
                    ".",
                    "--target",
                    "install",
                    "--config",
                    "Release",
                ],
                check=True,
                cwd=cmake_build_dir,
                capture_output=True,
            )
        except subprocess.CalledProcessError as e:
            self.app.display_error(f"cmake build failed: {e.stderr.decode()}")
            sys.exit(1)

        self.app.display_success("Base QIS build completed successfully")

    def build_platform_qis(self, platform: str):
        self.app.display_mini_header(f"Building QIS for {platform}")
        platform_qis_dir = Path(self.root) / f"selene-ext/interfaces/{platform}_qis"
        cmake_source_dir = platform_qis_dir / "c"
        cmake_build_dir = Path(self.root) / f"target/{platform}_qis_build"
        cmake_build_dir.mkdir(parents=True, exist_ok=True)
        dist_dir = platform_qis_dir / f"python/selene_{platform}_qis_plugin/_dist"
        dist_dir.mkdir(parents=True, exist_ok=True)
        selene_sim_dist_dir = Path(self.root) / "selene-sim/python/selene_sim/_dist"
        base_qis_dist_dir = (
            Path(self.root)
            / "selene-ext/interfaces/base_qis/python/selene_base_qis_plugin/_dist"
        )
        # when running in the source directory, we're building to:
        local_dist_lib_dir = dist_dir / "lib"
        # and we need to give it an rpath to point to selene's lib directory
        local_selene_dist_lib_dir = selene_sim_dist_dir / "lib"
        # so the rpath for local runs should be
        local_relative_path = os.path.relpath(
            local_selene_dist_lib_dir.resolve(), local_dist_lib_dir.resolve()
        )
        origin = self.runtime_origin()
        local_rpath = (
            origin if local_relative_path == "." else f"{origin}/{local_relative_path}"
        )

        local_base_qis_relative_path = os.path.relpath(
            (base_qis_dist_dir / "lib").resolve(), local_dist_lib_dir.resolve()
        )
        local_base_qis_rpath = f"{origin}/{local_base_qis_relative_path}"

        # but when running in an installed python environment, then
        # selene_sim and selene_platform_qis_plugin are actually siblings within
        # site-packages, so we need the path from
        # $site-packages/selene_platform_qis_plugin/_dist/lib
        # to
        # $site-packages/selene_sim/_dist/_lib
        # which is
        installed_selene_rpath = f"{origin}/../../../selene_sim/_dist/lib"
        installed_base_qis_rpath = f"{origin}/../../../selene_base_qis_plugin/_dist/lib"
        cmake_configure_cmd = [
            "cmake",
            "-DCMAKE_INSTALL_LIBDIR=lib",
            f"-DCMAKE_INSTALL_PREFIX={dist_dir}",
            *self.install_rpath_arg(
                [
                    local_rpath,
                    local_base_qis_rpath,
                    installed_selene_rpath,
                    installed_base_qis_rpath,
                ]
            ),
            "-DCMAKE_BUILD_TYPE=Release",
            f"-DCMAKE_PREFIX_PATH={selene_sim_dist_dir};{base_qis_dist_dir}",
            f"{cmake_source_dir}",
        ]
        if self.target_is_windows_gnu():
            cmake_configure_cmd = [
                cmake_configure_cmd[0],
                "-G",
                "MinGW Makefiles",
                *cmake_configure_cmd[1:],
            ]

        try:
            self.app.display_info(
                f"Running cmake configure: {' '.join(cmake_configure_cmd)}"
            )
            subprocess.run(
                cmake_configure_cmd,
                cwd=cmake_build_dir,
                check=True,
                capture_output=True,
            )
        except subprocess.CalledProcessError as e:
            if b"is different than the directory" not in e.stderr:
                self.app.display_error(f"cmake failed: {e.stderr.decode()}")
                sys.exit(1)
            try:
                # existing build dir is incompatible, delete and retry
                shutil.rmtree(cmake_build_dir)
                cmake_build_dir.mkdir()
                subprocess.run(
                    cmake_configure_cmd,
                    cwd=cmake_build_dir,
                    check=True,
                    capture_output=True,
                )
            except subprocess.CalledProcessError as e:
                self.app.display_error(f"cmake failed: {e.stderr.decode()}")
                sys.exit(1)

        try:
            subprocess.run(
                [
                    "cmake",
                    "--build",
                    ".",
                    "--target",
                    "install",
                    "--config",
                    "Release",
                ],
                check=True,
                cwd=cmake_build_dir,
                capture_output=True,
            )
        except subprocess.CalledProcessError as e:
            self.app.display_error(f"cmake build failed: {e.stderr.decode()}")
            sys.exit(1)

        self.app.display_success(f"QIS build for {platform} completed successfully")