More work on setup

Author: mikaem

mpi4py_fft/fftw/factory.py

@@ -4,6 +4,20 @@
 from .utilities import FFTW_FORWARD, FFTW_MEASURE
 
 def get_fftw_lib(dtype):
+    """Return compiled fftw module interfacing the FFTW library
+
+    Parameters
+    ----------
+    dtype : dtype
+        Data precision
+
+    Returns
+    -------
+    Module or ``None``
+        Module can be either :mod:`.fftwf_xfftn`, :mod:`.fftw_xfftn` or
+        :mod:`.fftwl_xfftn`, depending on precision.
+    """
+
     dtype = np.dtype(dtype).char.upper()
     if dtype == 'G':
         try:

mpi4py_fft/fftw/fftw_xfftn.pyx

@@ -292,4 +292,3 @@ cdef class FFT:
             output_array *= self._M
 
         return output_array
-

setup.py

@@ -3,6 +3,7 @@
 import os, sys
 import shutil
 from distutils import ccompiler
+import six
 from setuptools import setup
 from setuptools.extension import Extension
 from numpy import get_include
@@ -13,36 +14,27 @@
 include_dirs = [get_include(), os.path.join(sys.prefix, 'include')]
 library_dirs = [os.path.join(sys.prefix, 'lib')]
 for f in ('FFTW_ROOT', 'FFTW_DIR'):
-    if f in os.environ['PATH']:
+    if f in os.environ:
         library_dirs.append(os.path.join(os.environ[f], 'lib'))
         include_dirs.append(os.path.join(os.environ[f], 'include'))
 
+prec_map = {'float': 'f', 'double': '', 'long double': 'l'}
 compiler = ccompiler.new_compiler()
-assert compiler.find_library_file(library_dirs, 'fftw3'), 'Cannot find FFTW library!'
-has_threads = compiler.find_library_file(library_dirs, 'fftw3_threads')
 
-prec_map = {'float': 'fftwf_', 'double': 'fftw_', 'long double': 'fftwl_'}
-
-libs = {
-    'float': ['fftw3f'],
-    'double': ['fftw3'],
-    'long double': ['fftw3l']
-    }
-
-has_prec = ['double']
-for d in ('float', 'long double'):
-    if compiler.find_library_file(library_dirs, libs[d][0]):
-        has_prec.append(d)
-
-if has_threads:
-    for d in has_prec:
-        libs[d].append('_'.join((libs[d][0], 'threads')))
+libs = {}
+for d in ('float', 'double', 'long double'):
+    lib = 'fftw3'+prec_map[d]
+    if compiler.find_library_file(library_dirs, lib):
+        libs[d] = [lib]
+        tlib = '_'.join((lib, 'threads'))
+        if compiler.find_library_file(library_dirs, tlib):
+            libs[d].append(tlib)
         if sys.platform in ('unix', 'darwin'):
             libs[d].append('m')
 
 # Generate files with float and long double if needed
-for d in has_prec[1:]:
-    p = prec_map[d]
+for d in ('float', 'long double'):
+    p = 'fftw'+prec_map[d]+'_'
     for fl in ('fftw_planxfftn.h', 'fftw_planxfftn.c', 'fftw_xfftn.pyx', 'fftw_xfftn.pxd'):
         fp = fl.replace('fftw_', p)
         shutil.copy(os.path.join(fftwdir, fl), os.path.join(fftwdir, fp))
@@ -55,13 +47,13 @@
                  include_dirs=[get_include(),
                                os.path.join(sys.prefix, 'include')])]
 
-for d in has_prec:
-    p = prec_map[d]
+for d, v in six.iteritems(libs):
+    p = 'fftw'+prec_map[d]+'_'
     ext.append(Extension("mpi4py_fft.fftw.{}xfftn".format(p),
                          sources=[os.path.join(fftwdir, "{}xfftn.pyx".format(p)),
                                   os.path.join(fftwdir, "{}planxfftn.c".format(p))],
                          #define_macros=[('NPY_NO_DEPRECATED_API', 'NPY_1_7_API_VERSION')],
-                         libraries=libs[d],
+                         libraries=v,
                          include_dirs=include_dirs,
                          library_dirs=library_dirs))
 
@@ -92,5 +84,5 @@
           ],
       ext_modules=ext,
       install_requires=["mpi4py", "numpy", "six"],
-      setup_requires=["setuptools>=18.0", "cython>=0.25"]
+      setup_requires=["setuptools>=18.0", "cython>=0.25", "six"]
       )