better looking build

neuro-ml · Jun 4, 2024 · 46b5a6c · 46b5a6c
1 parent 688f4be
commit 46b5a6c
Showing 1 changed file with 30 additions and 21 deletions.
diff --git a/_build_utils.py b/_build_utils.py
@@ -63,34 +63,16 @@ def get_ext_modules():
         '/O3' if on_windows else '-O3',
     ]  # FIXME: account for higher gcc versions
 
-    # Cython extension and .pyx source file names must be the same to compile
-    # https://stackoverflow.com/questions/8024805/cython-compiled-c-extension-importerror-dynamic-module-does-not-define-init-fu
     modules = ['backprojection', 'measure', 'morphology', 'numeric', 'radon', 'zoom']
     modules_to_link_against_numpy_core_math_lib = ['numeric']
 
     src_dir = Path(__file__).parent / name / 'src'
     for module in modules:
+        # Cython extension and .pyx source file names must be the same to compile
+        # https://stackoverflow.com/questions/8024805/cython-compiled-c-extension-importerror-dynamic-module-does-not-define-init-fu
         shutil.copyfile(src_dir / f'_{module}.pyx', src_dir / f'_fast_{module}.pyx')
 
-    return [
-        Extension(
-            f'{name}.src._{prefix}{module}',
-            [f'{name}/src/_{prefix}{module}.pyx'],
-            include_dirs=[LazyImport('numpy')],
-            library_dirs=[NumpyLibImport()] if module in modules_to_link_against_numpy_core_math_lib else [],
-            libraries=(
-                ['npymath'] + ['m'] * (not on_windows) if module in modules_to_link_against_numpy_core_math_lib else []
-            ),
-            extra_compile_args=args + additional_args,
-            extra_link_args=args + additional_args,
-            define_macros=[('NPY_NO_DEPRECATED_API', 'NPY_1_7_API_VERSION')],
-        )
-        for module in modules
-        # FIXME: import of `ffast-math` compiled modules changes global FPU state, so now `fast=True` will just
-        # fallback to standard `-O2` compiled versions until https://github.com/neuro-ml/imops/issues/37 is resolved
-        # for prefix, additional_args in zip(['', 'fast_'], [[], ['-ffast-math']])
-        for prefix, additional_args in zip(['', 'fast_'], [[], []])
-    ] + [
+    extensions = [
         Extension(
             f'{name}.cpp.cpp_modules',
             [f'{name}/cpp/main.cpp'],
@@ -100,3 +82,30 @@ def get_ext_modules():
             language='c++',
         )
     ]
+    for module in modules:
+        libraries = []
+        library_dirs = []
+        include_dirs = [LazyImport('numpy')]
+
+        if module in modules_to_link_against_numpy_core_math_lib:
+            library_dirs.append(NumpyLibImport())
+            libraries.append('npymath')
+            if not on_windows:
+                libraries.append('m')
+
+        # FIXME: import of `ffast-math` compiled modules changes global FPU state, so now `fast=True` will just
+        # fallback to standard `-O2` compiled versions until https://github.com/neuro-ml/imops/issues/37 is resolved
+        # for prefix, additional_args in zip(['', 'fast_'], [[], ['-ffast-math']])
+        for prefix, additional_args in zip(['', 'fast_'], [[], []]):
+            extensions.append(
+                Extension(
+                    f'{name}.src._{prefix}{module}',
+                    [f'{name}/src/_{prefix}{module}.pyx'],
+                    include_dirs=include_dirs,
+                    library_dirs=library_dirs,
+                    libraries=libraries,
+                    extra_compile_args=args + additional_args,
+                    extra_link_args=args + additional_args,
+                    define_macros=[('NPY_NO_DEPRECATED_API', 'NPY_1_7_API_VERSION')],
+                )
+            )