How to register a destructor for a C-allocated numpy array?

后端 未结 1 1358
野趣味
野趣味 2021-02-09 02:03

I want to allocate numbers for a numpy array in C/C++, and pass them to python as a numpy array. That I can do with PyArray_SimpleNewFromData.

The problem is that I als

1条回答
  •  醉梦人生
    2021-02-09 02:34

    The idea is to create a Python object that knows how to free your memory when destroyed and make it the base of the returned C-allocated numpy array. This sounds tricky but it can be easily achieved via something known as capsules in python. Let me give an example,

    Suppose that you have the following code,

    PyObject *arr;
    int nd = 2;
    npy_intp dims[] = {5, 10};
    double *data = some_function_that_returns_a_double_star(x, y, z);
    
    arr = PyArray_SimpleNewFromData(nd, dims, NPY_DOUBLE, (void *)data);
    return arr;
    

    There is an obvious memory leak here since you cannot free data until arr is deleted as it says here in the red warning box. Fixing this, on the other hand, is easy. Define a function which is basically a destructor function that knows how to do garbage collection.

    void capsule_cleanup(PyObject *capsule) {
        void *memory = PyCapsule_GetPointer(capsule, NULL);
        // Use your specific gc implementation in place of free if you have to
        free(memory);
    }
    

    Now augment your code as,

    PyObject *arr;
    int nd = 2;
    npy_intp dims[] = {5, 10};
    double *data = some_function_that_returns_a_double_star(x, y, z);
    
    arr = PyArray_SimpleNewFromData(nd, dims, NPY_DOUBLE, (void *)data);
    PyObject *capsule = PyCapsule_New(data, NULL, capsule_cleanup);
    // NULL can be a string but use the same string while calling PyCapsule_GetPointer inside capsule_cleanup
    PyArray_SetBaseObject((PyArrayObject *) arr, capsule);
    return arr;
    

    There is no need to Py_DECREF the capsule. The function PyArray_SetBaseObject steals reference.

    Hope this helps!

    0 讨论(0)
提交回复
热议问题