setup.py

# Copyright python-stratify contributors
#
# This file is part of python-stratify and is released under the BSD license.
# See LICENSE in the root of the repository for full licensing details.
import os
from pathlib import Path
import sys
import warnings

# safe to import numpy here thanks to pep518
import numpy as np
from setuptools import Command, Extension, setup

try:
 from Cython.Build import cythonize # isort:skip
except ImportError:
 wmsg = "Cython unavailable, unable to build stratify extensions!"
 warnings.warn(wmsg)
 cythonize = None


BASE_DIR = Path(__file__).resolve().parent
PACKAGE_NAME = "stratify"
SRC_DIR = BASE_DIR / "src"
STRATIFY_DIR = SRC_DIR / PACKAGE_NAME
CMDS_NOCYTHONIZE = ["clean", "sdist"]
FLAG_COVERAGE = "--cython-coverage" # custom flag enabling Cython line tracing


class CleanCython(Command):
 description = "Purge artifacts built by Cython"
 user_options = []

 def initialize_options(self):
 pass

 def finalize_options(self):
 pass

 def run(self):
 for path in STRATIFY_DIR.rglob("*"):
 if path.suffix in (".pyc", ".pyo", ".c", ".so"):
 msg = f"clean: removing file {path}"
 print(msg)
 path.unlink()


cython_coverage_enabled = (
 os.environ.get("CYTHON_COVERAGE", None) or FLAG_COVERAGE in sys.argv
)
cython_directives = {}
define_macros = []
extensions = []

if cythonize and cython_coverage_enabled:
 define_macros.extend(
 [
 ("CYTHON_TRACE", "1"),
 ("CYTHON_TRACE_NOGIL", "1"),
 ]
 )
 cython_directives.update({"linetrace": True})
 if FLAG_COVERAGE in sys.argv:
 sys.argv.remove(FLAG_COVERAGE)
 print('enabled "linetrace" Cython compiler directive')

for fname in SRC_DIR.glob(f"{PACKAGE_NAME}/*.pyx"):
 # ref: https://setuptools.pypa.io/en/latest/userguide/ext_modules.html
 extension = Extension(
 f"{PACKAGE_NAME}.{fname.stem}",
 sources=[str(fname.relative_to(BASE_DIR))],
 include_dirs=[np.get_include()],
 define_macros=define_macros,
 )
 extensions.append(extension)

if cythonize and not any([arg in CMDS_NOCYTHONIZE for arg in sys.argv]):
 extensions = cythonize(
 extensions, compiler_directives=cython_directives, language_level=3
 )

cmdclass = {"clean_cython": CleanCython}
kwargs = {"cmdclass": cmdclass, "ext_modules": extensions}
setup(**kwargs)