numpy_demo

#include <Python.h>
#define NPY_NO_DEPRECATED_API NPY_1_7_API_VERSION
#include <numpy/arrayobject.h>

#include <iostream>

static PyObject* create_zero_array(PyObject* self, PyObject* args);

static PyMethodDef NpMethods[] = {
    {"create_zero_array", create_zero_array, METH_VARARGS, "create zero array of specified size"},
    {NULL, NULL, 0, NULL}
};

static PyModuleDef nptest_module = {
    PyModuleDef_HEAD_INIT,
    "nptest_module",
    "module for testing numpy",
    -1,
    NpMethods
};

PyMODINIT_FUNC PyInit_libnptest() {
    auto module = PyModule_Create(&nptest_module);
    import_array();
    return module;
}

static PyObject* create_zero_array(PyObject* self, PyObject* args) {
    int array_size;
    if (!PyArg_ParseTuple(args, "i", &array_size)) {
	return NULL;
    }

    npy_intp dims[1] = {array_size};
    double* data = new double[array_size];

    if (!data) {
	throw std::runtime_error("Unable to allocate storage");
    }

    PyObject* array = PyArray_SimpleNewFromData(1, dims, NPY_DOUBLE, data);

    for (size_t i = 0; i < array_size; ++i) {
	data[i] = 0;
    }
    PyArray_ENABLEFLAGS((PyArrayObject*) array, NPY_ARRAY_OWNDATA);

    return array;
}