尝试在 Cython 中为 cimport 声明（与 mpi4py 一起使用）

Question

有没有办法在 Cython 中为 cimport 提供等效的 Python try 语句？

类似的东西：

try:
    cimport something
except ImportError:
    pass

我需要这个来编写一个可以使用或不使用 mpi4py 编译的 Cython 扩展。这在编译语言中是非常标准的，其中 mpi 命令可以放在 #ifdef 和 #endif 预处理器指令之间。我们如何在 Cython 中获得相同的结果？

我试过了，但它不起作用：

try:
    from mpi4py import MPI
    from mpi4py cimport MPI
    from mpi4py.mpi_c cimport *
except ImportError:
    rank = 0
    nb_proc = 1

# solve a incompatibility between openmpi and mpi4py versions
cdef extern from 'mpi-compat.h': pass

does_it_work = 'Not yet'

实际上，如果 mpi4py 安装正确，但如果 import mpi4py 引发 ImportError，Cython 文件没有 编译，我得到错误：

Error compiling Cython file:
------------------------------------------------------------
...

try:
    from mpi4py import MPI
    from mpi4py cimport MPI
   ^
------------------------------------------------------------

mod.pyx:4:4: 'mpi4py.pxd' not found

文件setup.py：

from setuptools import setup, Extension
from Cython.Distutils import build_ext

import os
here = os.path.abspath(os.path.dirname(__file__))

include_dirs = [here]

try:
    import mpi4py
except ImportError:
    pass
else:
    INCLUDE_MPI = '/usr/lib/openmpi/include'
    include_dirs.extend([
        INCLUDE_MPI,
        mpi4py.get_include()])

name = 'mod'
ext = Extension(
    name,
    include_dirs=include_dirs,
    sources=['mod.pyx'])

setup(name=name,
      cmdclass={"build_ext": build_ext},
      ext_modules=[ext])

Answer 1

以这种方式使用 try-catch 块是您无法做到的。 您正在制作的扩展模块必须静态编译并链接到它使用 cimport 在 C 级别加载的内容。 try-catch 块是在导入模块时执行的，而不是在编译时执行的。

另一方面，理论上，您应该能够使用 Cython 的support for conditional compilation 获得您想要的效果。 在您的 setup.py 文件中，您可以检查是否定义了所需的模块，然后定义要传递给 Cython 编译器的环境变量，这又取决于所需的模块是否存在。

one of Cython's tests 中有一个如何执行此操作的示例。 在那里，他们将包含所需环境变量的字典传递给 Cython 的 Extension 类的构造函数，作为关键字参数 pyrex_compile_time_env，它一直是 renamed 到 cython_compile_time_env，而对于 Cython.Build.Dependencies.cythonize，则称为 compile_time_env）。

Answer 2

感谢您非常有用的回答@IanH。我举了一个例子来展示它的作用。

文件setup.py：

from setuptools import setup
from Cython.Distutils.extension import Extension
from Cython.Distutils import build_ext

import os
here = os.path.abspath(os.path.dirname(__file__))

import numpy as np
include_dirs = [here, np.get_include()]

try:
    import mpi4py
except ImportError:
    MPI4PY = False
else:
    MPI4PY = True
    INCLUDE_MPI = '/usr/lib/openmpi/include'
    include_dirs.extend([
        INCLUDE_MPI,
        mpi4py.get_include()])

name = 'mod'
ext = Extension(
    name,
    include_dirs=include_dirs,
    cython_compile_time_env={'MPI4PY': MPI4PY},
    sources=['mod.pyx'])

setup(name=name,
      cmdclass={"build_ext": build_ext},
      ext_modules=[ext])

if not MPI4PY:
    print('Warning: since importing mpi4py raises an ImportError,\n'
          '         the extensions are compiled without mpi and \n'
          '         will work only in sequencial.')

还有文件mod.pyx，带有一点真实的mpi命令：

import numpy as np
cimport numpy as np

try:
    from mpi4py import MPI
except ImportError:
    nb_proc = 1
    rank = 0
else:
    comm = MPI.COMM_WORLD
    nb_proc = comm.size
    rank = comm.Get_rank()

IF MPI4PY:
    from mpi4py cimport MPI
    from mpi4py.mpi_c cimport *

    # solve an incompatibility between openmpi and mpi4py versions
    cdef extern from 'mpi-compat.h': pass

    print('mpi4py ok')
ELSE:
    print('no mpi4py')

n = 8
if n % nb_proc != 0:
    raise ValueError('The number of processes is incorrect.')

if rank == 0:
    data_seq = np.ones([n], dtype=np.int32)
    s_seq = data_seq.sum()
else:
    data_seq = np.zeros([n], dtype=np.int32)

if nb_proc > 1:
    data_local = np.zeros([n/nb_proc], dtype=np.int32)
    comm.Scatter(data_seq, data_local, root=0)
else:
    data_local = data_seq

s = data_local.sum()
if nb_proc > 1:
    s = comm.allreduce(s, op=MPI.SUM)

if rank == 0:
    print('s: {}; s_seq: {}'.format(s, s_seq))
    assert s == s_seq

使用python setup.py build_ext --inplace 构建并使用python -c "import mod" 和mpirun -np 4 python -c "import mod" 进行测试。如果没有安装mpi4py，仍然可以构建模块并按顺序使用。