setup.py

"""Build hook for PyneCore's optional compiled extensions.

All package metadata lives in ``pyproject.toml``; this file declares the
extensions: ``pynecore.core._native_math``, the compiled twin of the pure-Python
transcendental functions in ``core/pine_math.py`` and ``core/fdlibm.py``, and
``pynecore.core._native_rolling_sum``, the compiled twin of the ``math.sum`` machine in
``core/rolling_sum.py``, and ``pynecore.core._native_color``, the compiled twin of
``lib.color.new``, and ``pynecore.core._native_profiler``, the ``sys.monitoring``
callbacks of the per-line profiler. ``pynecore.transformers._native`` is the Rust port
of the AST transform pipeline (``rust/``), built with setuptools-rust when a Rust
toolchain is present.

The extensions are optional in every sense: a platform without a C compiler, or a
compile that fails, still installs a working package that runs the pure-Python
implementations, and ``PYNE_BUILD_PURE=1`` builds the pure wheel on purpose (the one
WebAssembly runtimes and unlisted platforms install). The profiler is the exception:
it has no pure-Python twin, so such an install runs scripts but cannot profile them.
"""
import os
import shutil
import sys
from pathlib import Path

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

# The compiled code must round every operation exactly as the Python originals
# do. clang and gcc may fuse ``a * b + c`` into one FMA instruction unless told not to
# (clang does by default, and arm64 always has FMA), which changes last bits; fast-math
# would additionally reorder sums. MSVC contracts only under /fp:fast or /fp:contract.
_UNIX_FP_FLAGS = ['-ffp-contract=off', '-fno-fast-math', '-fno-unsafe-math-optimizations']
_MSVC_FP_FLAGS = ['/fp:precise']


class OptionalBuildExt(build_ext):
    """Compile with exact floating-point flags; on failure, fall back to pure Python."""

    def build_extensions(self) -> None:
        flags = _MSVC_FP_FLAGS if self.compiler.compiler_type == 'msvc' else _UNIX_FP_FLAGS
        for ext in self.extensions:
            ext.extra_compile_args = [*ext.extra_compile_args, *flags]
        super().build_extensions()

    def run(self) -> None:
        try:
            super().run()
        except Exception as exc:  # noqa: BLE001 -- any build failure means "pure Python"
            self._warn(exc)

    def build_extension(self, ext: Extension) -> None:
        try:
            super().build_extension(ext)
        except Exception as exc:  # noqa: BLE001 -- any build failure means "pure Python"
            self._warn(exc)

    @staticmethod
    def _warn(exc: Exception) -> None:
        print(f"WARNING: a pynecore native extension was not built ({exc}); "
              f"the pure-Python implementation will be used.", file=sys.stderr)


class BuildPyWithDocs(build_py):
    """Ship the repository's ``docs/`` inside the package as ``pynecore/docs/``.

    The documentation stays at the repository root; only the build copies it, so an
    installed PyneCore carries the docs of its own version (``pynecore.docs_path()``).
    An editable install skips the copy: there ``docs_path()`` resolves to the source
    tree's ``docs/``.
    """

    def run(self) -> None:
        super().run()
        if self.editable_mode:
            return
        source = Path(__file__).parent / 'docs'
        if not source.is_dir():
            raise FileNotFoundError(f"PyneCore documentation not found at {source}")
        target = Path(self.build_lib) / 'pynecore' / 'docs'
        if target.exists():
            shutil.rmtree(target)
        shutil.copytree(source, target)


def _extensions() -> list[Extension]:
    if os.environ.get('PYNE_BUILD_PURE'):
        return []
    try:
        from Cython.Build import cythonize
    except ImportError:
        print("WARNING: Cython not available; building without the native extensions.",
              file=sys.stderr)
        return []
    return cythonize(
        [Extension('pynecore.core._native_math', ['src/pynecore/core/_native_math.pyx']),
         Extension('pynecore.core._native_rolling_sum',
                   ['src/pynecore/core/_native_rolling_sum.pyx']),
         Extension('pynecore.core._native_color', ['src/pynecore/core/_native_color.pyx']),
         Extension('pynecore.core._native_profiler',
                   ['src/pynecore/core/_native_profiler.pyx'])],
        compiler_directives={'language_level': 3},
    )


def _rust_extensions() -> list:
    if os.environ.get('PYNE_BUILD_PURE'):
        return []
    try:
        from setuptools_rust import Binding, RustExtension
    except ImportError:
        print("WARNING: setuptools-rust not available; building without the native transform.",
              file=sys.stderr)
        return []
    # ``optional``: without a Rust toolchain, or when the build fails, the package
    # installs without it and every transform step runs in Python
    return [RustExtension('pynecore.transformers._native',
                          path='rust/crates/pyne-transform-py/Cargo.toml',
                          binding=Binding.PyO3, optional=True)]


_rust = _rust_extensions()
setup(ext_modules=_extensions(),
      cmdclass={'build_ext': OptionalBuildExt, 'build_py': BuildPyWithDocs},
      **({'rust_extensions': _rust} if _rust else {}))