Interfacing numpy and C: an example

#define NPY_NO_DEPRECATED_API NPY_API_VERSION

#include <Python.h>
#include <numpy/arrayobject.h>
#include <signal.h>

const char foo_docstring[] =
    "This function returns the (i,j) element of the matrix\n"
    "\n"
    "i,j are given in kwargs; each is assumed to be 0 if omitted.\n"
    "Matrix must have ndims == 2 and (i,j) must be in-bounds (exception\n"
    "thrown otherwise). The returned element is converted to a double-\n"
    "precision float, so the input matrix does not need to have dtype=float\n";
static PyObject* foo(PyObject* NPY_UNUSED(self),
                     PyObject* args,
                     PyObject* kwargs)
{
    char* keywords[] = {"arr",

                        // optional kwargs
                        "i", "j",
                        NULL};

    PyObject* result = NULL;

    // Python is silly. There's some nuance about signal handling where it sets
    // a SIGINT (ctrl-c) handler to just set a flag, and the python layer then
    // reads this flag and does the thing. Here I'm running C code, so SIGINT
    // would set a flag, but not quit, so I can't interrupt the C code as it's
    // running, which is very surprising to the user. Thus I reset the SIGINT
    // handler to the default, and put it back to the python-specific version
    // when I'm done. This ignores a potential custom python sigint handler.
    // Left as an excercise for the reader
    struct sigaction sigaction_old;
    if( 0 != sigaction(SIGINT,
                       &(struct sigaction){ .sa_handler = SIG_DFL },
                       &sigaction_old) )
    {
        PyErr_SetString(PyExc_RuntimeError, "sigaction() failed");
        goto done;
    }

    PyArrayObject* arr;
    unsigned int i = 0, j = 0; // defaults, if arguments not given

    if(!PyArg_ParseTupleAndKeywords( args, kwargs,
                                     "O&|II",
                                     keywords,
                                     PyArray_Converter, &arr,
                                     &i, &j))
        return NULL;

    {
        int        ndim     = PyArray_NDIM(arr);
        npy_intp*  dims     = PyArray_DIMS(arr);
        int        typenum  = PyArray_TYPE(arr);
        if( ndim != 2 )
        {
            PyErr_SetString(PyExc_RuntimeError, "I assume that ndims == 2");
            goto done;
        }
        if( i >= dims[0] )
        {
            PyErr_Format(PyExc_RuntimeError, "I assume that i < shape[0], but i == %d", i);
            goto done;
        }
        if( j >= dims[1] )
        {
            PyErr_Format(PyExc_RuntimeError, "I assume that j < shape[1], but j == %d", j);
            goto done;
        }

        if( typenum != NPY_DOUBLE )
        {
            PyArrayObject* arr2 =
                (PyArrayObject*)PyArray_FromArray(arr, PyArray_DescrFromType(NPY_DOUBLE),
                                                  NPY_ARRAY_ALIGNED);
            Py_DECREF(arr);
            arr = arr2;
        }
    }


    // Useful metadata about this matrix
    __attribute__((unused)) char*      data0    = PyArray_DATA    (arr);
    __attribute__((unused)) char*      data1    = PyArray_BYTES   (arr);
    __attribute__((unused)) npy_intp  *strides  = PyArray_STRIDES (arr);
    __attribute__((unused)) int        ndim     = PyArray_NDIM    (arr);
    __attribute__((unused)) npy_intp*  dims     = PyArray_DIMS    (arr);
    __attribute__((unused)) npy_intp   itemsize = PyArray_ITEMSIZE(arr);
    __attribute__((unused)) int        typenum  = PyArray_TYPE    (arr);


    // Two ways to grab the data out of the matrix:
    //
    // 1. Call a function. Clear what this does, but if we need to access a lot
    // of data in a loop, this has overhead: this function is pre-compiled, so
    // the overhead will not be inlined
    double d0 = *(double*)PyArray_GetPtr( arr, (npy_intp[]){i, j} );

    // 2. inline the function ourselves. It's pretty simple
    double d1 = *(double*)&data0[ i*strides[0] + j*strides[1] ];

    // The two methods should be identical. If not, this example is wrong, and I
    // barf
    if( d0 != d1 )
    {
        PyErr_Format(PyExc_RuntimeError, "PyArray_GetPtr() inlining didn't work: %f != %f", d0, d1);
        goto done;
    }


    result = Py_BuildValue( "d", d0 );

 done:
    Py_DECREF(arr);

    if( 0 != sigaction(SIGINT,
                       &sigaction_old, NULL ))
        PyErr_SetString(PyExc_RuntimeError, "sigaction-restore failed");
    return result;
}

PyMODINIT_FUNC inittst(void)
{
    static PyMethodDef methods[] =
        { {"foo", (PyCFunction)foo, METH_VARARGS | METH_KEYWORDS, foo_docstring},
         {}
        };


    PyImport_AddModule("tst");
    Py_InitModule3("tst", methods,
                   "Demo module to show binding numpy to C");

    // Required to avoid mysterious segfaults
    import_array();
}

#!/usr/bin/python

from setuptools import setup
from distutils.core import Extension

setup(name         = 'tst',
      version      = '0.1',
      author       = 'Dima Kogan',
      author_email = 'dima@secretsauce.net',
      ext_modules  = [Extension('tst',
                                sources = ['tst.c'])])

#!/usr/bin/python2

import numpy as np
import numpysane as nps
import sys

sys.path[:0] = ('/tmp/build/lib.linux-x86_64-2.7',)
import tst

x0 = np.arange(12).reshape(3,4)
x1 = nps.mv(x0, 0,1) # same matrix, with the axis order reversed

print tst.foo(x0, i=1, j=0)
print tst.foo(x1, i=0, j=1)

int main(int argc, char **argv)
{
    Py_SetProgramName(argv[0]);
    Py_Initialize();

    inittst();

    PySys_SetArgvEx(argc, argv, 0);

    PyRun_SimpleString("import numpy as np\n");
    PyRun_SimpleString("import numpysane as nps\n");
    PyRun_SimpleString("import tst\n");

    PyRun_SimpleString("x0 = np.arange(12).reshape(3,4)\n");
    PyRun_SimpleString("x1 = nps.mv(x0, 0,1)\n");
    PyRun_SimpleString("print tst.foo(x0, i=1, j=0)\n");
    PyRun_SimpleString("print tst.foo(x1, i=0, j=1)\n");

    Py_Exit(0);
    return 0;
}