C function to fully consume a Python iterator

Question

I was reading into the proper way to consume an iterator on this question as well as here and I notices there is no actual proper function for this that clearly explains itself.

Using:

collections.deque(maxlen=0).extend(iterable)

isn't exactly very descriptive whereas making a properly named method for this that just uses the deque approach is good I wanted to do my own research as well as learn some c (which I've rarely ever actually dealt with) along the way.

So.. Using the approach that is used in _collectionsmodule.c I have come up with this:

#include "Python.h"
#include <ctype.h>

static PyObject *
consume(PyObject *self, PyObject *iterable)
{
    PyObject *(*iternext)(PyObject *);
    PyObject *item;

    if (PyObject_GetIter(iterable) == NULL)
        return NULL;

    iternext = *Py_TYPE(iterable)->tp_iternext;
    while ((item = iternext(iterable)) != NULL)
        Py_DECREF(item);

    if (PyErr_Occurred()) {
        if (PyErr_ExceptionMatches(PyExc_StopIteration))
            PyErr_Clear();
        else
            return NULL;
    }
    Py_RETURN_NONE;
}

static PyMethodDef consumeMethods[] = {
    {"consume", (PyCFunction)consume, METH_O,
      "Run an iterator to exhaustion."},
    {NULL, NULL, 0, NULL}
};

static struct PyModuleDef _consumemodule = {
    PyModuleDef_HEAD_INIT,
    "_consumemodule",
    "Consume Module",
    -1,
    consumeMethods,
    NULL,
    NULL,
    NULL,
    NULL
};

PyMODINIT_FUNC
PyInit__consumemodule(void)
{
    return PyModule_Create(&_consumemodule);
}

It performs marginally better than the deque approach which is expected as it's using almost the same approach although I was wondering if any of the c guru's out there could see if I'm missing anything or if something could be done better as all it's doing it exhausting the iterator. I tried using METH_FASTCALL but that doesn't perform as well as METH_O and in the while loop I tried usingwhile ((Py_XDECREF(iternext(iterable))) != NULL) the only way I could find to eliminate the need for the variable item was by making a macro like this:

static inline int _Ret_DECREF(PyObject *op)
{
    if (op == NULL)
        return 0;
    Py_DECREF(op);
    return 1;
}

#define Ret_DECREF(op) _Ret_DECREF(op)

Then just looping through like this:

while (1) {
    if (Ret_DECREF(iternext(iterable)))
        break;
}

But something just feels off doing that.

Answer 1

You’ve got an object leak.

PyObject_GetIter() returns a new instance, which you aren’t calling Py_DECREF() on.

If you pass in an iterable object to your function (a list, a tuple, a dictionary, etc), PyObject_GetIter() will return a brand new iterator over than object. You test the return value, discover it is not NULL, and then lose the returned value, since it was not stored in a local variable. Then, you try to retrieve tp_iternext from the iterable’s type, which is likely not defined, so chaos will ensue.

If you pass in an iterator object to your function, PyObject_GetIter() will return it, but also increment the reference count of that object by one. At the end of your function, you don’t call Py_DECREF() on the returned iterator, so the reference count will remain incremented, and the iterator will never be released.

static PyObject *
consume(PyObject *self, PyObject *iterable)
{
    PyObject *(*iternext)(PyObject *);
    PyObject *item;
    PyObject *iterator;

    iterator = PyObject_GetIter(iterable);  // No references to iterable beyond this point.
    if (iterator == NULL)
        return NULL;

    iternext = *Py_TYPE(iterator)->tp_iternext;
    while ((item = iternext(iterator)) != NULL) 
        Py_DECREF(item);

    Py_DECREF(iterator);

    if (PyErr_Occurred()) {
        if (PyErr_ExceptionMatches(PyExc_StopIteration))
            PyErr_Clear();
        else
            return NULL;
    }

    Py_RETURN_NONE;
}