使用Cython为Python开发C++扩展

Cython的出现免去了为Python开发C/C++扩展的很多麻烦。本文以一个简单的例子来说明如何为Python开发C++扩展。

例子程序：给定一个列表，把列表的每个元素平方，并返回新列表。用Python实现会是这样：

def square(l):
    return [x * x for x in l]
1
2

现在我们用C++实现这个函数。根据Using C++ in Cython，Python列表对应于C++的std::vector，因此我们可以用std::vector。

_square.h:

#ifndef _SQUARE_H_
#define _SQUARE_H_

#include <vector>

std::vector<double> _square(std::vector<double> &);

#endif
1
2
3
4
5
6
7
8

_square.cpp:

#include "_square.h"

std::vector<double> _square(std::vector<double> &l)
{
    std::vector<double> res(l.size());
    for (auto i = l.begin(); i != l.end(); ++i) {
        res.push_back(*i * *i);
    }
    return res;
}
1
2
3
4
5
6
7
8
9
10

注意到上文代码文件名和函数名都以下划线开头，这里没有什么特殊规则，只是不让它们与Cython文件名和函数重名。接下来我们写封装C++的Cython代码。Cython代码后缀是.pyx。

square.pyx:

from libcpp.vector cimport vector

cdef extern from "_square.h":
    vector[double] _square(vector[double] l)

def square(l):
    cdef vector[double] l_vec = l
    return _square(l_vec)
1
2
3
4
5
6
7
8

最后我们编写用于编译的setup.py，setup.py位于项目根目录。这里假设上述_square.h、_square.cpp、square.pyx都位于Python package package1.package2下。

setup.py:

from setuptools import Extension, setup
from Cython.Build import cythonize

extensions = [
    Extension(
        # 这里写完整包名
        name='package1.package2.square',
        # 这里包含Cython文件和C++源文件
        sources=[
            'package1/package2/square.pyx',
            'package1/package2/_square.cpp',
        ],
        # 这里写编译flags；
        # - 写`-std=c++11`因为我们用了`auto`关键字
        # - 写`-DNDEBUG`是为了忽略所有`assert`（虽然这里并没有`assert`，只是为多举一个例子）
        extra_compile_args=['-std=c++11', '-DNDEBUG'],
        language='c++',
    ),
]

setup(
    # name参数可写可不写，这里没写
    #name='...',
    ext_modules=cythonize(extensions),
    zip_safe=False,
)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26

注意最后有一个zip_safe=False，根据Building a Cython module using setuptools，这是为避免一个导入错误：

One caveat: the default action when running python setup.py install is to create a zipped egg file which will not work with cimport for pxd files when you try to use them from a dependent package. To prevent this, include zip_safe=False in the arguments to setup().

最后我们来编译这个扩展模块。在命令行，项目根目录（即setup.py所在目录），执行：

python3 setup.py build_ext --inplace
1

为执行这条命令，Windows需要Visual Studio，Linux需要GNU工具链（g++），Mac需要XCode（clang++)。

为使用这个扩展模块，我们可以这样：

from package1.package2.square import square

l1 = [1., 2., 3.]
print(square(l1))
1
2
3
4

输出

[1.0, 4.0, 9.0]
1

致谢

本文受这个回答启发而创作。