setup.py

from __future__ import annotations

import platform

import numpy
import setuptools
from Cython.Build import cythonize
from setuptools.command.build_ext import build_ext
from setuptools.errors import CCompilerError
from setuptools_rust import Binding, RustExtension

ext_modules = cythonize(
    module_list=[
        setuptools.Extension(
            "*",
            sources=["river/**/*.pyx"],
            include_dirs=[numpy.get_include()],
            libraries=[] if platform.system() == "Windows" else ["m"],
            define_macros=[("NPY_NO_DEPRECATED_API", "NPY_1_7_API_VERSION")],
        )
    ],
    compiler_directives={
        "binding": True,
        "embedsignature": True,
    },
)

rust_extensions = [RustExtension("river.stats._rust_stats", binding=Binding.PyO3)]


class BuildFailed(Exception):
    pass


class ExtBuilder(build_ext):
    def run(self):
        try:
            build_ext.run(self)
        except FileNotFoundError:
            raise BuildFailed("File not found. Could not compile C extension.")

    def build_extension(self, ext):
        try:
            build_ext.build_extension(self, ext)
        except (CCompilerError, ValueError):
            raise BuildFailed("Could not compile C extension.")


setuptools.setup(
    ext_modules=ext_modules,
    rust_extensions=rust_extensions,
    cmdclass={"build_ext": ExtBuilder},
    zip_safe=False,
    include_package_data=True,
)