setup.py

import sys
import os
import shutil
import pathlib
import multiprocessing

from math import ceil

import toml

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


def get_project_name() -> str:
    with open(pathlib.Path().absolute() / "pyproject.toml", "r") as file:
        project_toml = toml.load(file)
        return project_toml["project"]["name"]


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


class CMakeBuild(build_ext):
    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))

        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)

        os.chdir(str(build_temp))

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

        build_gen = (
            ["-G", os.environ["AIDGE_BUILD_GEN"]]
            if "AIDGE_BUILD_GEN" in os.environ
            else []
        )

        self.spawn(
            [
                "cmake",
                *build_gen,
                str(cwd),
                "-DTEST=OFF",
                f"-DCMAKE_INSTALL_PREFIX:PATH={install_path}",
                f"-DCMAKE_BUILD_TYPE={compile_type}",
                "-DPYBIND=ON",
                "-DCMAKE_EXPORT_COMPILE_COMMANDS=ON",
                "-DCOVERAGE=OFF",
            ]
        )

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

        aidge_package = build_lib / (get_project_name())

        # Get "aidge core" package
        # ext_lib = build_temp
        print(build_temp.absolute())
        # Copy all shared object files from build_temp/lib to aidge_package
        for root, _, files in os.walk(build_temp.absolute()):
            for file in files:
                if (file.endswith(".so") or file.endswith(".pyd")) and (
                    root != str(aidge_package.absolute())
                ):
                    currentFile = os.path.join(root, file)
                    shutil.copy(currentFile, str(aidge_package.absolute()))


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