Search code examples
pythonc++memory-managementpython-extensions

Understanding memory leak with C++ extension for python


I have been struggling to understand what I am doing wrong with the memory management of this this C++ function for a python module, but I don't have much experience in this regard.

Every time I run this function, the memory consumption of the python interpreter increases. The function is supposed to take in two numpy arrays and create new numpy arrays with the desired output.

extern "C" PyObject *radec_to_xyz(PyObject *self, PyObject *args) {
    // Parse the input arguments
    PyArrayObject *ra_arrobj, *dec_arrobj;
    if (!PyArg_ParseTuple(args, "O!O!", &PyArray_Type, &ra_arrobj, &PyArray_Type, &dec_arrobj)) {
        PyErr_SetString(PyExc_TypeError, "invalid arguments, expected two numpy arrays");
        return nullptr;
    }
    // skipping checks that would ensure:
    // dtype==float64, dim==1, len()>0 and equal for inputs, data or contiguous
    npy_intp size = PyArray_SIZE(ra_arrobj);

    // create the output numpy array with the same size and datatype
    PyObject *x_obj = PyArray_EMPTY(1, &size, NPY_FLOAT64, 0);
    if (!x_obj) return nullptr;
    Py_XINCREF(x_obj);

    // get pointers to the arrays
    double *ra_array = static_cast<double *>(PyArray_DATA(ra_arrobj));
    double *dec_array = static_cast<double *>(PyArray_DATA(dec_arrobj));
    double *x_array = static_cast<double *>(PyArray_DATA(reinterpret_cast<PyArrayObject*>(x_obj)));

    // compute the new coordinates
    for (npy_intp i = 0; i < size; ++i) {
        double cos_ra = cos(ra_array[i]);
        double cos_dec = cos(dec_array[i]);
        // compute final coordinates
        x_array[i] = cos_ra * cos_dec;
    }
    // return the arrays holding the new coordinates
    return Py_BuildValue("O", x_obj);
}

I suspect that I am getting the reference counts wrong and therefore, the returned numpy arrays don't get garbage collected. I tried changing the reference counts, but this did not help.

When I pass the X array as additional argument from the python interpreter instead of allocating it in the function, the memory leak is gone as expected.


Solution

  • Py_XINCREF(x_obj) is not needed because Py_BuildValue("O", x_obj) will increment the reference count for you. In other words, you're accidentally increasing the reference count too much by 1