setup.py

import sys
import os
import shutil
import pathlib
import multiprocessing

from math import ceil

from setuptools import setup, Extension
from setuptools.command.build_ext import build_ext


def get_project_name() -> str:
    return open(pathlib.Path().absolute() / "project_name.txt", "r").read().strip()


PROJECT_NAME = get_project_name()

SETUP_DIR = pathlib.Path(__file__).parent

class CMakeExtension(Extension):
    def __init__(self, name):
        super().__init__(name, sources=[])


class CMakeBuild(build_ext):
    def __init__(self, dist, *args, **kwargs):
        super().__init__(dist, *args, **kwargs)
        # Detect editable_mode for old versions of setuptools
        if not hasattr(self, "editable_mode"):
            if hasattr(dist, "commands"):
                self.editable_mode = "develop" in dist.commands
            else:
                self.editable_mode = False

    def run(self):
        # This lists the number of processors available on the machine
        # The compilation will use half of them
        max_jobs = str(ceil(multiprocessing.cpu_count() / 2))
        max_jobs = os.environ.get("AIDGE_NB_PROC", max_jobs)

        cwd = pathlib.Path().absolute()

        build_temp = cwd / "build"
        if not build_temp.exists():
            build_temp.mkdir(parents=True, exist_ok=True)

        build_lib = pathlib.Path(self.build_lib)
        if not build_lib.exists():
            build_lib.mkdir(parents=True, exist_ok=True)

        package_prefix = build_lib if not self.editable_mode else SETUP_DIR
        pybind_install_prefix = (package_prefix / PROJECT_NAME).absolute()

        install_path = (
            os.path.join(sys.prefix, "lib", "libAidge")
            if "AIDGE_INSTALL" not in os.environ
            else os.environ["AIDGE_INSTALL"]
        )

        # Read environment variables for CMake options
        c_compiler = os.environ.get("AIDGE_C_COMPILER", "gcc")
        cxx_compiler = os.environ.get("AIDGE_CXX_COMPILER", "g++")
        build_type = os.environ.get("AIDGE_BUILD_TYPE", "Release")
        asan = os.environ.get("AIDGE_ASAN", "OFF")
        cmake_arch = os.environ.get("AIDGE_CMAKE_ARCH", "")

        build_gen = os.environ.get("AIDGE_BUILD_GEN", "")
        build_gen_opts = (
            ["-G", build_gen]
            if build_gen
            else []
        )
        test_onoff = os.environ.get("AIDGE_BUILD_TEST", "OFF")

        os.chdir(str(build_temp))

        cmake_cmd = [
            "cmake",
            *build_gen_opts,
            str(cwd),
            f"-DTEST={test_onoff}",
            f"-DCMAKE_INSTALL_PREFIX:PATH={install_path}",
            f"-DCMAKE_BUILD_TYPE={build_type}",
            f"-DCMAKE_C_COMPILER={c_compiler}",
            f"-DCMAKE_CXX_COMPILER={cxx_compiler}",
            f"-DENABLE_ASAN={asan}",
            "-DPYBIND=ON",
            f"-DPYBIND_INSTALL_PREFIX:PATH={pybind_install_prefix}",
            "-DCMAKE_EXPORT_COMPILE_COMMANDS=1",
            "-DCOVERAGE=OFF",
        ]

        # Append architecture-specific arguments if provided
        if cmake_arch:
            cmake_cmd.append(cmake_arch)

        self.spawn(cmake_cmd)

        if not self.dry_run:
            self.spawn(
                ["cmake", "--build", ".", "--config", build_type, "-j", max_jobs]
            )
            self.spawn(["cmake", "--install", ".", "--config", build_type])
        os.chdir(str(cwd))


if __name__ == "__main__":
    setup(
        ext_modules=[CMakeExtension(PROJECT_NAME)],
        cmdclass={
            "build_ext": CMakeBuild,
        },
        zip_safe=False,
    )