用C++写一个调用numpy的Python模块

目前有一些分析的项目要用到numpy,但某些操作numpy并没有提供对应的函数,如果用python代码来编写这些功能的话在数据量较大的时候时间开销就不能忽略了,因此这篇文章说下用C++写一个Python扩展,并与numpy交互的一个例子。

本文实现的功能很简单:对于一个二维的numpy多维int类型数组,对每一行,我们需要输出其大于某个值的个数。例如用Python实现的话代码如下:

import numpy as np

def large_than_target(data, n):
    result = dict()
    for i in range(0, data.shape[0]):
        result[i] = 0
        for j in range(0, data.shape[1]):
            if data[i][j] > n:
                result[i] += 1
    return [result[k] for k in result]

data = np.array([[1, 2, 3], [5, 6, 7]])
print(large_than_target(data, 2))

输出为:

[1, 3]

我们先来建立一个目录,名字叫做pyextnp,然后建立下面几个文件:

(python3.5)➜  pyextnp tree
.
├── pyextnp
│   └── __init__.py
├── setup.py
└── src
    └── pyextnp.cpp

2 directories, 5 files

setup.py的内容如下:

from distutils.core import setup, Extension

import numpy as np

EXTRA_COMPILE_ARGS = ['-std=c++11', '-stdlib=libc++', '-mmacosx-version-min=10.8']
EXTRA_LINK_ARGS = ['-std=c++11', '-stdlib=libc++', '-mmacosx-version-min=10.8']

pyextnpcpp_module = Extension('pyextnpcpp',
                              sources=['src/pyextnp.cpp'],
                              include_dirs=[np.get_include()],
                              extra_compile_args=EXTRA_COMPILE_ARGS,
                              extra_link_args=EXTRA_LINK_ARGS)

REQUIRES = []

setup(name='PyExtNp',
      version='1.0',
      description='PyExtNp',
      author='Qin TianHuan',
      author_email='tianhuan@bingotree.cn',
      packages=['pyextnp'],
      ext_modules=[pyextnpcpp_module],
      requires=REQUIRES)

这个setup.py和普通的纯python代码的区别就在于这里用到了Extension。实际上最终生成package的时候,我们的C++的代码是被编译成一个动态链接库的,python解释器在加载C++代码时会加载对应的动态链接库,而这里的Extension则会指示具体编译时要用到哪些源文件,同时制定对应的编译、链接选项。关于这一块的详细说明可以参考distutils的文档。现在我们来build一下:

(python3.5)➜  pyextnp python setup.py build
running build
running build_py
creating build
creating build/lib.macosx-10.6-intel-3.5
creating build/lib.macosx-10.6-intel-3.5/pyextnp
copying pyextnp/__init__.py -> build/lib.macosx-10.6-intel-3.5/pyextnp
running build_ext
building 'pyextnpcpp' extension
creating build/temp.macosx-10.6-intel-3.5
creating build/temp.macosx-10.6-intel-3.5/src
/usr/bin/clang -fno-strict-aliasing -Wsign-compare -Wunreachable-code -fno-common -dynamic -DNDEBUG -g -fwrapv -O3 -Wall -Wstrict-prototypes -arch i386 -arch x86_64 -g -I/Users/thuanqin/Desktop/Dev/lab/python3.5/lib/python3.5/site-packages/numpy/core/include -I/Library/Frameworks/Python.framework/Versions/3.5/include/python3.5m -c src/pyextnp.cpp -o build/temp.macosx-10.6-intel-3.5/src/pyextnp.o -std=c++11 -stdlib=libc++ -mmacosx-version-min=10.8
/usr/bin/clang++ -bundle -undefined dynamic_lookup -arch i386 -arch x86_64 -g build/temp.macosx-10.6-intel-3.5/src/pyextnp.o -o build/lib.macosx-10.6-intel-3.5/pyextnpcpp.cpython-35m-darwin.so -std=c++11 -stdlib=libc++ -mmacosx-version-min=10.8

可以看到,distutils工具会自动为我们编译并生成pyextnpcpp.cpython-35m-darwin.so库。我们install下:

(python3.5)➜  pyextnp python setup.py install
running install
running build
running build_py
running build_ext
running install_lib
creating /Users/thuanqin/Desktop/Dev/lab/python3.5/lib/python3.5/site-packages/pyextnp
copying build/lib.macosx-10.6-intel-3.5/pyextnp/__init__.py -> /Users/thuanqin/Desktop/Dev/lab/python3.5/lib/python3.5/site-packages/pyextnp
copying build/lib.macosx-10.6-intel-3.5/pyextnpcpp.cpython-35m-darwin.so -> /Users/thuanqin/Desktop/Dev/lab/python3.5/lib/python3.5/site-packages
byte-compiling /Users/thuanqin/Desktop/Dev/lab/python3.5/lib/python3.5/site-packages/pyextnp/__init__.py to __init__.cpython-35.pyc
running install_egg_info
Writing /Users/thuanqin/Desktop/Dev/lab/python3.5/lib/python3.5/site-packages/PyExtNp-1.0-py3.5.egg-info

现在import下我们的package,注意此时我们有一个可以import的python package和一个C++的python module:

>>> import pyextnp
>>> import pyextnpcpp
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
ImportError: dynamic module does not define module export function (PyInit_pyextnpcpp)

可以看到pyextnpcpp已经可以被解释器识别了,但是由于这个so文件没有对外暴露PyInit_pyextnpcpp符号因此解释器无法正常加载这个so文件。下面我们就来编写我们的pyextnp.cpp。

首先,基本的框架代码如下:

#include <iostream>
#include <vector>
#include <map>

#include "numpy_import.h"
#ifdef NO_IMPORT_ARRAY
#undef NO_IMPORT_ARRAY
#endif
#include "numpy/arrayobject.h"
#include "Python.h"

static PyObject *pyextnp_error;

static PyObject *
large_than_target(PyObject *self, PyObject *args)
{
    Py_RETURN_NONE;
}

static PyMethodDef PyExtNpMethods[] = {
        {"large_than_target",  large_than_target, METH_VARARGS, "large_than_target"},
        {NULL, NULL, 0, NULL}
};

static struct PyModuleDef pyextnpmodule = {
        PyModuleDef_HEAD_INIT,
        "pyextnp",
        NULL,
        -1,
        PyExtNpMethods
};

PyMODINIT_FUNC
PyInit_pyextnpcpp(void)
{
    PyObject *m;

    m = PyModule_Create(&pyextnpmodule);
    if (m == NULL)
        return NULL;

    pyextnp_error = PyErr_NewException("pyextnp.error", NULL, NULL);
    Py_INCREF(pyextnp_error);
    PyModule_AddObject(m, "error", pyextnp_error);

    import_array();

    return m;
}

这里的numpy_import.h内容为:

#ifndef KAGE_COMMON_H
#define KAGE_COMMON_H

#define NO_IMPORT_ARRAY
#define PY_ARRAY_UNIQUE_SYMBOL PYEXTNP_ARRAY_API

#endif //KAGE_COMMON_H

此时我们python setup.py install下后就能调用了:

(python3.5)➜  pyextnp python                 
Python 3.5.1 (v3.5.1:37a07cee5969, Dec  5 2015, 21:12:44) 
[GCC 4.2.1 (Apple Inc. build 5666) (dot 3)] on darwin
Type "help", "copyright", "credits" or "license" for more information.
>>> import pyextnpcpp
>>> pyextnpcpp.large_than_target()
>>> 

要理解这些代码只要了解python解释器的调用逻辑就行了,具体的顺序如下:
1.执行import pyextnpcpp,解释器去环境变量指定的路径下搜索对应的so文件然后用dlopen加载,并定位so文件的PyInit_XXX符号,XXX就是我们的模块名
2.找到PyInit_pyextnpcpp后执行该函数,该函数被PyMODINIT_FUNC宏包裹,PyMODINIT_FUNC其实是一个extern “C”用于兼容C++编译器
3.PyInit_pyextnpcpp首先调用PyModule_Create创建本module,而其参数pyextnpmodule则包含了这个module的相关信息,比较重要的是其PyMethodDef *这个成员指针,用于指定这个module对外暴露哪些方法,例如我们这里的large_than_target方法。这样在python中import了pyextnpcpp模块后就能调用large_than_target方法了。解释器在发现Python代码调用large_than_target后,就会去调用so文件的large_than_target这个C++方法了。
4.PyErr_NewException用于定义一个异常,这里异常的一般反馈方法是用PyErr_SetString给pyextnp_error设置一个报错信息然后返回NULL
5.PyModule_AddObject用于在我们的module中添加这个异常对象
6.import_array这个用于让python解释器import下numpy相关的依赖。因为我们的这个例子中我们的C++代码需要依赖另一个Python的C扩展numpy,如果运行过程中Python解释器还没加载对应的numpy扩展,那么我们的C++代码在执行时就会出现问题。import_array和我们numpy_import.h头文件中的PY_ARRAY_UNIQUE_SYMBOL宏相关,因为对于多个文件链接生成的且都使用了numpy相关函数的代码来说,import_array会生成一个类似于PY_ARRAY_UNIQUE_SYMBOL_XXX的符号并对外暴露,每一个使用了numpy功能的C++代码都应该使用它们自己这个扩展的PY_ARRAY_UNIQUE_SYMBOL_XXX,因此import_array在这里会为我们的扩展生成PYEXTNP_ARRAY_API_XXX这样的符号。

现在来实现我们的large_than_target函数,代码如下:

static PyObject *
large_than_target(PyObject *self, PyObject *args)
{
    PyArrayObject *array = nullptr;
    int n;
    std::map<int, int> result;
    if (!PyArg_ParseTuple(args, "O!i", &PyArray_Type, &array, &n)) {
        PyErr_SetString(pyextnp_error, "ARGS ERROR");
        return NULL;
    }

    if (array->nd != 2) {
        PyErr_SetString(pyextnp_error, "ND ERROR");
        return NULL;
    }

    for (auto i=0; i<array->dimensions[0]; i++) {
        result[i] = 0;
        for (auto j=0; j<array->dimensions[1]; j++) {
            auto value = PyArray_GETPTR2(array, i, j);
            if(*(int *)value > n) {
                result[i] += 1;
            }
        }
    }

    PyObject* list = PyList_New(0);
    for (auto it=result.begin(); it!=result.end(); it++) {
        PyList_Append(list, Py_BuildValue("i", it->second));
    }

    return list;
}

这里的代码比较简单,我们先通过PyArg_ParseTuple获取解释器传来的参数,由于在解释器中一切都是PyObject对象,因此我们的C++代码需要将其转换为对应的C++对象。然后我们通过array的dimensions遍历,这里的逻辑就和Python的代码是一样的了。

最后我们来看下性能差距吧,我们对每个调用都调用1000000次,结果如下:

(python3.5)➜  time python cpp.py 
python cpp.py  1.04s user 0.04s system 99% cpu 1.095 total
(python3.5)➜  time python py.py 
python py.py  7.03s user 0.07s system 99% cpu 7.132 total

最后,一般我们提供代码的时候是不会直接提供扩展模块的,而是用python代码做个包装,毕竟python简单许多。例如我们这边的例子,一般用户只要使用pyextnp即可,而pyextnp则会使用pyextnpcpp。

关于Python的C++模块编写,可以参考:
# 扩展基础
https://docs.python.org/2/extending/extending.html
# 如何扩展一个新的类别
https://docs.python.org/2/extending/newtypes.html
# 部署
https://docs.python.org/2/extending/building.html

关于numpy的C扩展,可以参考:
http://docs.scipy.org/doc/numpy/reference/c-api.html

发表评论

电子邮件地址不会被公开。 必填项已用*标注

*