setup.py

"""Build script for the sdpb-python Cython extension.

The extension links against static libraries produced by SDPB's own waf build
(``c-src/sdpb/build``).  Build SDPB first, e.g. with ``scripts/build_sdpb.sh``,
then ``pip install -e .``.

Environment variables:
    SDPB_BUILD_DIR   Path to the waf build directory (default: c-src/sdpb/build).
    SDPB_SRC_DIR     Path to the SDPB source tree (default: c-src/sdpb).
    CXX / CC         Compiler; defaults to ``mpicxx`` because SDPB requires MPI.
"""

import os
import re
import subprocess
from pathlib import Path

from Cython.Build import cythonize
from setuptools import Extension, setup

ROOT = Path(__file__).resolve().parent
SDPB_SRC = Path(os.environ.get("SDPB_SRC_DIR", ROOT / "c-src" / "sdpb")).resolve()
SDPB_BUILD = Path(os.environ.get("SDPB_BUILD_DIR", SDPB_SRC / "build")).resolve()
PKG = ROOT / "src" / "sdpb_python"

# SDPB is compiled with mpicxx; the extension must be too.
os.environ.setdefault("CXX", "mpicxx")
os.environ.setdefault("CC", "mpicxx")


def sdpb_version() -> str:
    try:
        return subprocess.check_output(
            ["git", "describe", "--tags", "--always", "--dirty"],
            cwd=SDPB_SRC, text=True,
        ).strip()
    except (OSError, subprocess.CalledProcessError):
        return "unknown"


def waf_cache() -> dict:
    """Read the flags waf discovered at configure time so we link the same way."""
    cache = SDPB_BUILD / "c4che" / "_cache.py"
    if not cache.exists():
        return {}
    env: dict = {}
    exec(cache.read_text(), env)  # waf writes plain python assignments
    return {k: v for k, v in env.items() if not k.startswith("__")}


# Packages that only SDPB's command-line tools and file-based SDP input use:
# MPSolve (GPL-3, `spectrum`), libxml2 (XML input), libarchive (sdp.zip).  The
# extension contains none of that code.  scripts/build_sdpb.sh configures SDPB
# with --libs-only, which does not even look for them; they are excluded here as
# well so that an extension built on top of a full SDPB build links the same way.
EXCLUDED_PACKAGES = {"mpsolve", "libxml2", "libarchive"}


def collect(env: dict, prefix: str) -> list:
    out: list = []
    for key, value in env.items():
        if key.startswith(prefix + "_") and isinstance(value, list):
            if key[len(prefix) + 1:] in EXCLUDED_PACKAGES:
                continue
            for item in value:
                if item not in out:
                    out.append(item)
    return out


env = waf_cache()

include_dirs = [
    str(PKG / "cpp"),
    str(SDPB_SRC / "src"),
    str(SDPB_SRC / "external"),
    *collect(env, "INCLUDES"),
]
library_dirs = [str(SDPB_BUILD), *collect(env, "LIBPATH"), *collect(env, "STLIBPATH")]
runtime_library_dirs = collect(env, "RPATH")


def parse_define(d: str) -> tuple:
    name, _, value = d.partition("=")
    return (name, value or None)


waf_defines = [parse_define(d) for d in collect(env, "DEFINES") + list(env.get("DEFINES", []))]

# Static libs built by waf (order matters for static linking) plus external libs.
sdpb_static_libs = ["sdp_solve", "pmp2sdp_lib", "pmp", "sdpb_util"]
libraries = sdpb_static_libs + collect(env, "LIB") + collect(env, "STLIB")
if not env:
    # Reasonable defaults when the waf cache is unavailable.
    libraries += ["El", "boost_program_options", "boost_serialization",
                  "gmpxx", "gmp", "mpfr", "flint", "openblas"]


def rel(path: Path) -> str:
    """setuptools requires source paths relative to setup.py."""
    return os.path.relpath(path, ROOT).replace(os.sep, "/")


extensions = [
    Extension(
        "sdpb_python._sdpb",
        sources=[
            rel(PKG / "_sdpb.pyx"),
            rel(PKG / "cpp" / "sdpb_wrapper.cxx"),
        ],
        include_dirs=include_dirs,
        library_dirs=library_dirs,
        runtime_library_dirs=runtime_library_dirs,
        libraries=libraries,
        language="c++",
        extra_compile_args=["-std=c++17", "-O3", "-Wall"],
        define_macros=[
            ("OMPI_SKIP_MPICXX", None),
            ("SDPB_VERSION_STRING", f'"{sdpb_version()}"'),
            *waf_defines,
        ],
    )
]

setup(
    ext_modules=cythonize(
        extensions,
        language_level=3,
        compiler_directives={"embedsignature": True},
    ),
)