Skip to content
Snippets Groups Projects
Code owners
Assign users and groups as approvers for specific file changes. Learn more.
setup.py 2.96 KiB
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))

        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()

        os.chdir(str(build_temp))

        compile_type = os.environ.get("AIDGE_PYTHON_BUILD_TYPE", "Release")
        install_path = (
            os.path.join(sys.prefix, "lib", "libAidge")
            if "AIDGE_INSTALL" not in os.environ
            else os.environ["AIDGE_INSTALL"]
        )
        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")

        self.spawn(
            [
                "cmake",
                *build_gen_opts,
                str(cwd),
                f"-DTEST={test_onoff}",
                f"-DCMAKE_INSTALL_PREFIX:PATH={install_path}",
                f"-DCMAKE_BUILD_TYPE={compile_type}",
                "-DPYBIND=ON",
                f"-DPYBIND_INSTALL_PREFIX:PATH={pybind_install_prefix}",
                "-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))


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