setup.py

from setuptools import setup, Extension
from setuptools.command.build import build as _build
from setuptools.command.build_ext import build_ext as _build_ext
from setuptools.command.build_py import build_py as _build_py
import numpy as np
import shutil
import os
import sys
import shlex

import platform

IS_WINDOWS = platform.system() == "Windows"
IS_DARWIN = platform.system() == "Darwin"
IS_LINUX = platform.system() == "Linux"

# When set, build a single-rank variant that requires neither a native MPI
# library nor mpi4py: bundled stub headers replace <mpi.h> and <mpi4py/mpi4py.h>,
# and the build records the choice in repast4py/_mpi_config.py so the Python side
# (repast4py._mpi) matches what the _space extension was compiled against.
SINGLE_RANK = os.environ.get("R4PY_NO_MPI", "") not in ("", "0")

SRC_DIR = os.path.join("src", "repast4py")
STUB_INCLUDE_DIR = os.path.join(SRC_DIR, "single_rank_include")


def is_64bit() -> bool:
    """Returns True if the python interpreter is 64-bit, independent of the OS arch."""
    return sys.maxsize > 2**32


def run_command(exe, args):
    cmd = shutil.which(exe)
    if not cmd:
        return []
    if not isinstance(args, str):
        args = " ".join(args)
    try:
        with os.popen(cmd + " " + args) as f:
            return shlex.split(f.read())
    except Exception:
        return []


def get_mpi_include_dirs():
    """MPI/mpi4py include directories.

    Single-rank builds use the bundled stub headers and do not depend on mpi4py.
    """
    if SINGLE_RANK:
        return [STUB_INCLUDE_DIR]
    try:
        import mpi4py
    except ImportError:
        # mpi4py is a *dynamic* build requirement (see build_support/r4py_backend.py).
        # During the PEP 517 get_requires / egg_info phase it may not be installed
        # yet; the include dir is only needed when the extensions are actually
        # compiled, by which point pip has installed mpi4py.
        return []
    return [mpi4py.get_include()]


def get_linker_args():
    linker_args = []

    # Single-rank builds link no MPI library.
    if SINGLE_RANK:
        return linker_args

    # NOTE Windows setuptools apparently does not use the CC env variable.
    if IS_WINDOWS:
        pass
    else:
        compiler = os.getenv("CC")
        if compiler is None:
            print(
                'Error: MPI compiler is not specified. Please specify the MPI compiler using the "CC" environment variable'
            )
            # Nothing to query for link flags. Commands that do not compile
            # (sdist, egg_info) still need setup.py to import.
            return linker_args
        args = run_command(compiler, "-show")
        for arg in args:
            if arg.startswith("-l") or arg.startswith("-L"):
                linker_args.append(arg)
        # linker_args_str = ' '.join(linker_args)
        # return linker_args_str
    return linker_args


def get_compiler_args():
    if IS_WINDOWS:
        compile_args = ["/std:c++latest"]
    else:
        compile_args = ["-std=c++11"]

    return compile_args


def get_extra_includes():
    # Single-rank builds need no MSMPI headers.
    if IS_WINDOWS and not SINGLE_RANK:
        return [os.environ["MSMPI_INC"]]
    else:
        return []


def get_lib_dirs():
    if IS_WINDOWS and not SINGLE_RANK:
        if is_64bit():
            return [os.environ["MSMPI_LIB64"]]
        else:
            return [os.environ["MSMPI_LIB32"]]
    else:
        return []


def get_libs():
    if IS_WINDOWS and not SINGLE_RANK:
        return ["msmpi"]
    else:
        return []


def write_mpi_config(package_dir):
    """Record the build choice so repast4py._mpi imports the matching MPI.

    Written into whichever copy of the package is about to be used: the build
    output for a wheel, the source tree for an in-place or editable build.
    """
    os.makedirs(package_dir, exist_ok=True)
    path = os.path.join(package_dir, "_mpi_config.py")
    with open(path, "w") as f:
        f.write("# Generated by setup.py - do not edit.\n")
        f.write("# Records whether repast4py was built single-rank (no native MPI / mpi4py).\n")
        f.write(f"SINGLE_RANK = {SINGLE_RANK}\n")


class build(_build):
    """Give each build mode its own build directory.

    The two modes compile the same sources against different headers, so they
    must not share object files: nothing about the sources changes when the mode
    does, and build_ext compares timestamps only. Sharing a directory lets a
    build in one mode silently relink the other mode's objects.
    """

    def finalize_options(self):
        if SINGLE_RANK and self.build_base == "build":
            self.build_base = os.path.join("build", "nompi")
        super().finalize_options()


class build_py(_build_py):
    """Write the marker into the build output, so it is packaged into the wheel
    without touching the source tree."""

    def run(self):
        super().run()
        if not self.dry_run:
            write_mpi_config(os.path.join(self.build_lib, "repast4py"))


class build_ext(_build_ext):
    """Write the marker into the source tree for in-place builds.

    `build_ext --inplace` puts the extensions next to the sources and is what
    developers run; editable installs reach here too, since setuptools turns
    `editable_mode` into `inplace`. Either way the package is imported from
    `src/`, so that is where the marker has to be.
    """

    def run(self):
        super().run()
        if self.inplace and not self.dry_run:
            write_mpi_config(SRC_DIR)


def get_install_requires():
    """Runtime dependencies. mpi4py is required only for the default (native MPI)
    build; the single-rank build uses the bundled pure-Python substitute."""
    requires = [
        "numpy",
        "numba",
        "torch",
        "networkx >=2.6.2",
        "pyyaml",
        "Cython",
        'typing_extensions;python_version < "3.8.0"',
        "dill",
    ]
    if not SINGLE_RANK:
        requires.append("mpi4py")
    return requires


# Single-rank stub headers must be found before any system MPI headers.
mpi_include_dirs = get_mpi_include_dirs()

core_module = Extension(
    "repast4py._core",
    sources=["src/repast4py/coremodule.cpp"],
    language="c++",
    extra_compile_args=get_compiler_args(),
    depends=["core.h"],
    extra_link_args=get_linker_args(),
    include_dirs=[np.get_include(), *mpi_include_dirs],
)
space_module = Extension(
    "repast4py._space",
    sources=[
        "src/repast4py/spacemodule.cpp",
        "src/repast4py/geometry.cpp",
        "src/repast4py/space.cpp",
        "src/repast4py/distributed_space.cpp",
        "src/repast4py/SpatialTree.cpp",
    ],
    language="c++",
    extra_compile_args=get_compiler_args(),
    depends=[
        "space.h",
        "grid.h",
        "cspace.h",
        "space_types.h",
        "geometry.h",
        "distributed_space.h",
        "borders.h",
    ],
    extra_link_args=get_linker_args(),
    include_dirs=[np.get_include(), *mpi_include_dirs, *get_extra_includes()],
    library_dirs=get_lib_dirs(),
    libraries=get_libs(),
)

setup(
    description="repast4py package",
    ext_modules=[core_module, space_module],
    install_requires=get_install_requires(),
    cmdclass={"build": build, "build_py": build_py, "build_ext": build_ext},
)