3
\$\begingroup\$

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.

\$\endgroup\$
2
\$\begingroup\$

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;
}
| improve this answer | |
\$\endgroup\$
  • \$\begingroup\$ Thanks for the great breakdown and helpful answer although, this is exactly the same as is used in the deque's consume_iterator and that may just be the best option but what I was thinking was to see if there is a better alternative that doesn’t need to store item for every iteration. I’ve tried do { Py_DECREF(iternext(iterator)) } while (!PyErr_Occurred()) but maybe my lack of c knowledge is making me think that there’s a better option while there isn’t. \$\endgroup\$ – Jab Jun 21 '19 at 13:32
  • 1
    \$\begingroup\$ @Jab If you look at consume_iterator, you’ll see it ends with return finalize_iterator(it);, and in that is the Py_DECREF(it);. I understand you were hoping for a speed improvement, not a bug fix, but there is a bug in your code and I’ve showed you how to fix it. The consume_iterator function doesn’t have the bug because it calls the additional helper function to clean up the iterator reference. You didn’t copy that part into your implementation. \$\endgroup\$ – AJNeufeld Jun 21 '19 at 13:58
  • \$\begingroup\$ I must have taken that out by mistake as I did look at those docs and saw it did make a new reference before even posting this too. Thanks for showing me \$\endgroup\$ – Jab Jun 21 '19 at 14:02
  • \$\begingroup\$ Your do { ... } while ( ) loop doesn’t work because iternext(iterator) will return NULL which will cause Py_DECREF(...) to crash. You would need to use Py_XDECREF(...) which allows the argument to be NULL. \$\endgroup\$ – AJNeufeld Jun 21 '19 at 14:04
  • \$\begingroup\$ I meant Py_XDECREF I’m on my phone away from the desk. Although when I tried that I believe it was an infinite loop or it only ran the first iteration I’ll have to try again when I get back to see. \$\endgroup\$ – Jab Jun 21 '19 at 14:08

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service, privacy policy and cookie policy

Not the answer you're looking for? Browse other questions tagged or ask your own question.