setup.py

#!/usr/bin/env python

import ctypes.util
import glob
import os
import re
import sys
import subprocess

from packaging import version as packaging_version
from setuptools import Extension, setup
from setuptools.command.build_ext import build_ext as _build_ext

import versioneer

# This is needed to use numpy in this module, and should work whether or not numpy is
# already installed. If it's not, it will trigger an installation

_default_xspec_version = "12.15.1"  # default when installing xspec according following
# https://heasarc.gsfc.nasa.gov/docs/software/conda.html


class My_build_ext(_build_ext):

    def finalize_options(self):

        _build_ext.finalize_options(self)

        import numpy

        self.include_dirs.append(numpy.get_include())
        self.include_dirs.append("astromodels/xspec/include")
        conda_prefix = os.environ.get("CONDA_PREFIX")
        if conda_prefix is not None:
            conda_include_path = os.path.join(conda_prefix, "include")
            self.include_dirs.append(conda_include_path)

    def run(self):
        """Build extensions and fix library references on macOS"""
        _build_ext.run(self)

        # On macOS, fix library references for libraries in extra_objects
        if sys.platform.lower().find("darwin") >= 0:
            self._fix_macos_library_references()

    def _fix_macos_library_references(self):
        """Fix install_name references for libraries without symlinks on macOS"""
        if not self.extensions:
            return

        for ext in self.extensions:
            if not hasattr(ext, "extra_objects") or not ext.extra_objects:
                continue

            # Get the output file path
            ext_path = self.get_ext_fullpath(ext.name)

            if not os.path.exists(ext_path):
                continue

            # Get library references using otool
            try:
                result = subprocess.run(
                    ["otool", "-L", ext_path],
                    capture_output=True,
                    text=True,
                    check=True,
                )
            except (subprocess.CalledProcessError, FileNotFoundError):
                continue

            # Build a map of expected broken references to actual library paths
            ref_map = {}
            for lib_path in ext.extra_objects:
                if not os.path.exists(lib_path):
                    continue

                lib_name = os.path.basename(lib_path)
                # Check if it's a versioned library (e.g., libwcs.8.3.dylib)
                match = re.match(r"lib(.+)\.(\d+)\.(\d+)\.dylib", lib_name)
                if match:
                    base_name, major, minor = match.groups()
                    # The library might be referenced as lib{base}.{major}.dylib
                    broken_ref = f"lib{base_name}.{major}.dylib"
                    # Use @rpath if library_dirs is set, otherwise use absolute path
                    if hasattr(ext, "library_dirs") and ext.library_dirs:
                        new_ref = f"@rpath/{lib_name}"
                    else:
                        new_ref = lib_path
                    ref_map[broken_ref] = new_ref

            # Parse otool output to find library references that need fixing
            for line in result.stdout.split("\n"):
                line = line.strip()
                if not line or ext_path in line:
                    continue

                # Extract library path (format: "libname.dylib (compatibility ...)")
                # After strip(), there's no leading whitespace
                match = re.match(r"^(.+?)\s+\(", line)
                if not match:
                    continue

                lib_ref = match.group(1)

                # Check if this reference needs to be fixed
                if lib_ref in ref_map:
                    new_ref = ref_map[lib_ref]
                    try:
                        subprocess.run(
                            [
                                "install_name_tool",
                                "-change",
                                lib_ref,
                                new_ref,
                                ext_path,
                            ],
                            check=True,
                            capture_output=True,
                        )
                        print(f"Fixed library reference: {lib_ref} -> {new_ref}")
                    except subprocess.CalledProcessError:
                        pass


def sanitize_lib_name(library_path):
    """Get a fully-qualified library name, like /usr/lib/libgfortran.so.3.0,
    and returns the lib name needed to be passed to the linker in the -l option
    (for example gfortran)

    :param library_path:
    :return:
    """

    lib_name = os.path.basename(library_path)

    # Some regexp magic needed to extract in a system-independent (mac/linux)
    # way the library name

    tokens = re.findall(r"lib(.+)(\.so|\.dylib|\.a|\.la)(.+)?", lib_name)

    if not tokens:
        msg = f"Attempting to find {lib_name} in directory {library_path}"
        msg += " but there are no libraries in this directory"

        raise RuntimeError(msg)

    return tokens[0][0]


def find_library(library_root, additional_places=None):
    """Returns the name of the library without extension.

    :param library_root: root of the library to search, for example
        "cfitsio_" will match libcfitsio_1.2.3.4.so
    :return: a tuple of (library_name, library_dir, full_path) where:
        - library_name: the name to be passed to the linker in the
          -l option
        - library_dir: the directory path (None if in system paths)
        - full_path: the full path to the library file (used when
          unversioned symlink is missing)
    """

    # find_library searches for all system paths in a system independent way (but NOT
    # those defined in LD_LIBRARY_PATH or DYLD_LIBRARY_PATH)

    first_guess = ctypes.util.find_library(library_root)

    if first_guess is not None and library_root == "gfortran":

        # Found in one of the system paths

        if sys.platform.lower().find("linux") >= 0:

            # On linux the linker already knows about these paths, so we
            # can return None as path

            return sanitize_lib_name(first_guess), None, None

        elif sys.platform.lower().find("darwin") >= 0:

            # On Mac we still need to return the path, because the linker sometimes
            # does not look into it

            return sanitize_lib_name(first_guess), os.path.dirname(first_guess), None

        else:

            # Windows is not supported

            raise NotImplementedError("Platform %s is not supported" % sys.platform)

    else:

        # could not find it. Let's examine LD_LIBRARY_PATH or DYLD_LIBRARY_PATH
        # (if they sanitize_lib_name(first_guess), are not defined, possible_locations
        # will become [""] which will be handled by the next loop)

        if sys.platform.lower().find("linux") >= 0:

            # Unix / linux

            possible_locations = os.environ.get("LD_LIBRARY_PATH", "").split(":")

        elif sys.platform.lower().find("darwin") >= 0:

            # Mac

            possible_locations = os.environ.get("DYLD_LIBRARY_PATH", "").split(":")

        else:

            raise NotImplementedError("Platform %s is not supported" % sys.platform)

        if additional_places is not None:

            possible_locations.extend(additional_places)

        # Now look into the search paths

        library_name = None
        library_dir = None
        library_full_path = None

        for search_path in possible_locations:

            if search_path == "":
                # This can happen if there are more than one :, or if nor
                # LD_LIBRARY_PATH nor DYLD_LIBRARY_PATH are defined (because of
                # the default use above for os.environ.get)

                continue

            results = glob.glob(os.path.join(search_path, f"lib{library_root}*"))

            if len(results) >= 1:

                # Results contain things like libXS.so, libXSPlot.so, libXSpippo.so
                # If we are looking for libXS.so, we need to make sure that we get the
                # right one!

                for result in results:

                    if (
                        re.match(
                            f"lib{library_root}" + r"[\-_\.]([0-9])*\d*(\.[0-9]\d*)*",
                            os.path.basename(result),
                        )
                        is None
                    ):

                        continue

                    else:

                        # FOUND IT

                        # This is the full path of the library, like
                        # /usr/lib/libcfitsio_1.2.3.4

                        library_full_path = result
                        library_name = result
                        library_dir = search_path

                        break

            else:

                continue

            if library_name is not None:
                break

        if library_name is None:

            return None, None, None

        else:

            # Sanitize the library name to get from the fully-qualified path
            # to just the library name (/usr/lib/libgfortran.so.3.0 becomes
            # gfortran)

            sanitized_name = sanitize_lib_name(library_name)

            # Extract base library name (without version suffix)
            # e.g., "wcs.8.3" -> "wcs", "gfortran" -> "gfortran"
            base_name = sanitized_name.split(".")[0]

            # Check if the unversioned symlink exists
            # If not, we'll need to pass the full path to the linker
            if sys.platform.lower().find("linux") >= 0:
                extension = ".so"
            elif sys.platform.lower().find("darwin") >= 0:
                extension = ".dylib"
            else:
                extension = ".so"

            # Check for unversioned symlink (e.g., libwcs.so or libwcs.dylib)
            unversioned_lib = os.path.join(library_dir, f"lib{base_name}{extension}")

            # On macOS, also check for major version symlink (e.g., libwcs.8.dylib)
            # which is what the runtime linker expects
            if sys.platform.lower().find("darwin") >= 0:
                # Extract major version if present (e.g., "8.3" -> "8")
                version_parts = sanitized_name.split(".")[1:]
                if version_parts:
                    major_version = version_parts[0]
                    major_version_lib = os.path.join(
                        library_dir, f"lib{base_name}.{major_version}{extension}"
                    )
                    if os.path.exists(major_version_lib):
                        # Major version symlink exists, use normal linking
                        return f"{base_name}.{major_version}", library_dir, None

            if os.path.exists(unversioned_lib):
                # Unversioned symlink exists, use normal linking
                return base_name, library_dir, None
            else:
                # No unversioned symlink, return the full path for direct
                # linking
                return sanitized_name, library_dir, library_full_path


def get_xspec_conda_version():
    """Get the version string from conda"""
    try:
        import xspec

        return xspec.Xset.version[1]
    except ModuleNotFoundError:
        return None


def setup_xspec():

    skip_xspec = os.environ.get("SKIP_XSPEC")
    headas_root = os.environ.get("HEADAS")
    conda_prefix = os.environ.get("CONDA_PREFIX")

    if skip_xspec is not None:

        print(
            "The SKIP_XSPEC env variable was set. Xspec support will not be installed."
        )
        return None

    if headas_root is None:

        # See, maybe we are running in Conda

        if conda_prefix is None:

            # Maybe this is a Conda build

            conda_prefix = os.environ.get("PREFIX")

        if conda_prefix is not None:

            # Yes, this is Conda
            # Let's see if the package xspec-modelsonly has been installed by checking
            # whether one of the Xspec libraries exists within conda
            conda_lib_path = os.path.join(conda_prefix, "lib")
            this_lib, this_lib_path, full_lib_library = find_library(
                "XSFunctions", additional_places=[conda_lib_path]
            )

            if this_lib is None:

                # No, there is no library in Conda
                msg = "No xspec-modelsonly package has been installed in Conda. Xspec"
                msg += " support will not be installed"
                print(msg)

                print("Was looking into %s" % conda_lib_path)

                return None

            else:
                msg = (
                    "WARN: The xspec-modelsonly package has been installed"
                    " in Conda, but it's no longer supported."
                    " Xspec support will not be installed"
                )
                print(msg)

                return None

                # Set up the HEADAS variable so that the following will find the
                # libraries
                # headas_root = conda_prefix

        else:

            print("No HEADAS env. variable set. Xspec support will not be installed ")

            return None

    print("HEADAS env. variable detected. Will compile the Xspec extension.")
    print(
        "NOTICE: If you have issues, manually set the environment variable "
        "XSPEC_INC_PATH to the location of the XSPEC headers"
    )
    msg = "If you are still having issues, unset HEADAS before installing and"
    msg += "contact the support team"
    print(msg)

    xspec_version = get_xspec_conda_version()

    if xspec_version is not None:

        print("Found XSPEC version %s in Conda" % xspec_version)

    else:

        print("No XSPEC installation found in Conda")
        print("Xspec was likely compiled from source.")

        xspec_version = os.environ.get("ASTRO_XSPEC_VERSION")

        if xspec_version is None:
            print("WARN: You have not specified an XSPEC version with the ")
            print("WARN: environment variable ASTRO_XSPEC_VERSION")
            print(f"WARN: we will assume you have {_default_xspec_version}")
            print(
                "If you are using a different version of XSPEC, please set"
                " the environment variable ASTRO_XSPEC_VERSION to the "
                "version of XSPEC you are using"
            )

            xspec_version = _default_xspec_version

        else:

            print(f"You have specified the XSPEC version {xspec_version}")

    xspec_version = packaging_version.Version(xspec_version)

    if xspec_version < packaging_version.Version("12.12.0"):
        msg = "WARN: XSPEC version is less than 12.12.0, which is the minimal"
        msg += " supported version for astromodels"
        print(msg)
        return None
    elif xspec_version > packaging_version.Version("12.15.1"):
        msg = "WARN: XSPEC version is greater than 12.15.1, which is the"
        msg += " maximal supported version for astromodels"
        print(msg)
        return None

    macros = []
    # I am not sure what the naming of the XSPEC components are,
    # but let's stick with major, minor, and patch.
    for major, minor, patch in [
        (12, 12, 0),
        (12, 12, 1),
        (12, 13, 0),
        (12, 13, 1),
        (12, 14, 0),
        (12, 14, 1),
        (12, 15, 0),
        (12, 15, 1),
    ]:

        version = "{}.{}.{}".format(major, minor, patch)

        macro = "XSPEC_{}_{}_{}".format(major, minor, patch)

        if xspec_version >= packaging_version.Version(version):
            macros += [(macro, None)]

    print(macros)

    # Make sure these libraries exist and are linkable right now
    # (they need to be in LD_LIBRARY_PATH or DYLD_LIBRARY_PATH or in one of the system
    # paths)

    libraries_root = [
        "XSFunctions",
        "XSModel",
        "XSUtil",
        "XS",
        "cfitsio",
        "CCfits",
        "wcs",
        "gfortran",
    ]

    libraries = []
    library_dirs = []
    extra_objects = []  # For libraries without unversioned symlinks

    for lib_root in libraries_root:

        this_library, this_library_path, full_lib_path = find_library(
            lib_root, additional_places=[os.path.join(headas_root, "lib")]
        )

        if this_library is None:

            raise IOError(
                "Could not find library %s. Impossible to compile Xspec" % lib_root
            )

        else:

            print("Found library %s in %s" % (this_library, this_library_path))

            if full_lib_path is not None:
                # No unversioned symlink exists, pass the full path directly
                print(
                    "Warning: No unversioned symlink found for %s, "
                    "using full path %s" % (lib_root, full_lib_path)
                )
                extra_objects.append(full_lib_path)
            else:
                # Normal case: unversioned symlink exists, use -l linking
                libraries.append(this_library)

            if this_library_path is not None:
                # This library is not in one of the system path library, we need to add
                # it to the -L flag during linking. Let's put it in the library_dirs
                # list which will be used in the Extension class

                library_dirs.append(this_library_path)

    # try to manually add on the include directory

    header_paths = []

    if library_dirs:

        # grab it from the lib assuming that it is one up
        xspec_path, _ = os.path.split(library_dirs[0])
        include_path = os.path.join(xspec_path, "include")

        header_paths.append(include_path)

    # let's be sure to add the conda include directory

    if conda_prefix is not None:

        conda_include_path = os.path.join(conda_prefix, "include")
        header_paths.append(conda_include_path)

    # check if there are user set the location of the xspec headers:

    xspec_headers_path = os.environ.get("XSPEC_INC_PATH")

    if xspec_headers_path is not None:

        print("You have set XSPEC_INC_PATH=%s" % xspec_headers_path)

        header_paths.append(xspec_headers_path)

    # Remove duplicates from library_dirs and header_paths

    library_dirs = list(set(library_dirs))
    header_paths = list(set(header_paths))

    print("header paths:")
    for h in header_paths:

        print(f"{h}")

    # Configure the variables to build the external module with the C/C++ wrapper

    ext_modules_configuration = [
        Extension(
            "astromodels.xspec._xspec",
            [
                "astromodels/xspec/src/_xspec.cc",
            ],
            include_dirs=header_paths,
            libraries=libraries,
            library_dirs=library_dirs,
            runtime_library_dirs=library_dirs,
            extra_compile_args=[],
            # Add full paths for libraries without unversioned symlinks
            extra_objects=extra_objects,
            define_macros=macros,
        ),
    ]

    return ext_modules_configuration


# Normal packages

packages = [
    "astromodels",
    "astromodels/core",
    "astromodels/functions",
    "astromodels/functions/functions_1D",
    "astromodels/functions/dark_matter",
    "astromodels/sources",
    "astromodels/utils",
    "astromodels/xspec",
    "astromodels/tests",
]

# Check whether we can compile Xspec support
ext_modules_configuration = setup_xspec()

# Add the node_ctype module


setup(
    cmdclass=versioneer.get_cmdclass({"build_ext": My_build_ext}),
    version=versioneer.get_version(),
    ext_modules=ext_modules_configuration,
)