Использование Numpy C api в C-расширении для Python

Ниже рассмотрен пример использования Numpy C api при написании расширения для Python на языке C. Расширение будет подключаться с помощью distutils. В примере умножаются поэлементно два массива.

Среда:

  • Linux
  • Python 3.5
  • IPython

Ниже C-код расширения и комментарии. Пояснения размещены под кодом.

#include 
#include 

#include 

/* Disable old Numpy API */
#define NPY_NO_DEPRECATED_API NPY_1_7_API_VERSION
#include <numpy/npy_math.h>
#include "numpy/arrayobject.h"

/* Errors */
static PyObject *MultiplyError;

/*
 * Enter point from Python
 */
static PyObject*
py_multiply(PyObject* self, PyObject* args)
{
    PyArrayObject *in_arr=NULL, *in_arr_second=NULL, *out_arr=NULL;
    int nd = 0, nd_second = 0;

    /* Parse arguments from Python:
       "O!" - check type of object
       &in_arr  - to variable in_arr
     */
    if (!PyArg_ParseTuple(args, "O!O!", &PyArray_Type, &in_arr, &PyArray_Type, &in_arr_second))
        return NULL;

    /* Get number of dimensions
     */
    nd = PyArray_NDIM(in_arr);
    nd_second = PyArray_NDIM(in_arr_second);

    /* Get length of array from shape
       [0] - length
     */
    int length = (int)PyArray_SHAPE(in_arr)[0];
    int length_second = (int)PyArray_SHAPE(in_arr_second)[0];

    /* Error with different sizes
     */
    if (length != length_second) {
        PyErr_SetString(MultiplyError, "Arrays have different sizes");
        return NULL;
    }
    /* Error with wrong ndim
     */
    if (nd != 1 || nd != nd_second) {
        PyErr_SetString(MultiplyError, "Arrays must have 1 dimension");
        return NULL;
    }

    /* Create array with zeros
     */
    out_arr = (PyArrayObject *) PyArray_ZEROS(nd, PyArray_SHAPE(in_arr), NPY_DOUBLE, 0);

    /* Get pointer to the first elements of arrays
     */
    double *item = (double *)PyArray_DATA(in_arr);
    double *item_second = (double *)PyArray_DATA(in_arr_second);
    double *item_out = (double *)PyArray_DATA(out_arr);

    double *end = (item + length);
    double *end_second = (item_second + length_second);
    for (int i = 0; item != end && item_second != end_second; item++, item_second++, i++) {
        item_out[i] = (*item) * (*item_second);
    }

    /* Return new array without increase reference count:
     * O - increase reference count
     * N - not increase reference count
     */
    return Py_BuildValue("N", out_arr);
}

/* Array with methods
 */
static PyMethodDef module_methods[] =
{
     /* name from python, name in C-file, ..., __doc__ string of method */
     {"multiply", py_multiply, METH_VARARGS, "Multiply two numpy arrays."},
     {NULL, NULL, 0, NULL}
};

/* Array with info about module to create it in Python
 */
static struct PyModuleDef moduledef =
{
    PyModuleDef_HEAD_INIT,
    "numpy_multiply",               /* name of module */
    "Multiply two numpy arrays in C-extensions",  /* module documentation, may be NULL */
    -1,                     /* size of per-interpreter state of the module,
                               or -1 if the module keeps state in global variables. */
    module_methods
};

/* Init our module in Python
 */
PyMODINIT_FUNC PyInit_numpy_multiply(void)
{
    PyObject *m;
    m = PyModule_Create(&moduledef);
    if (!m) {
        return NULL;
    }

    /* Import NUMPY settings
     */
    import_array();

    /* Init errors */
    MultiplyError = PyErr_NewException("multiply.error", NULL, NULL);
    Py_INCREF(MultiplyError); /* Increment reference count for object */
    PyModule_AddObject(m, "error", MultiplyError);

    return m;
}

файл numpy_multiply.c

В начале файла, по рекомендации разработчика Numpy, отключаем обратную совместимость со старой версией Numpy C api. Это позволит писать актуальный код и сократить проблемы в будущем:

#define NPY_NO_DEPRECATED_API NPY_1_7_API_VERSION

В начале функции py_multiply() получаем numpy-массивы автоматически проверяя тип полученных переменных инструкцией "O!" и на выходе имея переменную с типом PyArrayObject.

Ниже представлены примеры с бросанием исключения в Python при неверном измерении и разных размерах массивов.

Ближе к концу, создаем пустой массив размера, равного одному из входящих массивов и заполняем его нулями. При создании массива для полученного объекта счетчик внешних ссылок автоматически приравнивается 1.

Теперь в цикле умножаем элементы массивов. Этот участок кода можно реализовать как угодно в рамках правил языка C.

Возвращаем полученный массив с инструкцией "N", это означает, что при передаче объекта мы не будем увеличивать количество ссылок на него. В случае использования инструкции "O" можно легко получить утечку памяти, так как количество ссылок будет увеличено на 1 и этот счетчик уже никто не обнулит.

IPython

Все дальнейшие действия выполняются в IPython. Каждый кусок кода вставляем в отдельные ячейки.

Создаем python-скрипт setup_numpy_multiply.py для компиляции библиотеки в python-пакет.

%%writefile setup_numpy_multiply.py
from distutils.core import setup, Extension
import numpy as np
ext = Extension('numpy_multiply', 
                sources = ['numpy_multiply.c'])
setup(name="Numpy Mupltiply in C-Extension", 
      include_dirs = [np.get_include()], #Add Include path of numpy
      ext_modules = [ext]
     )

ipython, ячейка 1

Запускаем созданный на предыдущем шаге файл setup_numpy_multiply.py для компиляции в текущую директорию C-библиотеки.

%%bash
python setup_numpy_multiply.py build_ext --inplace

ipython, ячейка 2

Подключаем модуль и тестируем.

import sys
import numpy as np
import numpy_multiply

arr = np.linspace(0, 20, 21);
arr_second = np.linspace(0, 20, 21);
out = numpy_multiply.multiply(arr, arr_second)

arr, arr_second, out, sys.getrefcount(arr), sys.getrefcount(arr_second), sys.getrefcount(out)

ipython, ячейка 3

Репозиторий с кодом.