Improvements to coding for itertools.tee():
* Add error checking code to PyList_Append() call.
* Replace PyObject_CallMethod(to->outbasket, "pop", NULL) with equivalent
in-line code. Inlining is important here because the search for the
pop method will occur for every element returned by the iterator.
* Make tee's dealloc() a little smarter. If the trailing iterator is
being deallocated, then the queue data is no longer needed and can
be freed.
diff --git a/Modules/itertoolsmodule.c b/Modules/itertoolsmodule.c
index 42440df..b37dec1 100644
--- a/Modules/itertoolsmodule.c
+++ b/Modules/itertoolsmodule.c
@@ -46,6 +46,7 @@
{
teeobject *to = lz->tee;
PyObject *result, *tmp;
+ int n;
if (lz->num_seen == to->num_seen) {
/* This instance is leading, use iter to get more data */
@@ -53,7 +54,10 @@
if (result == NULL)
return NULL;
if (to->save_mode)
- PyList_Append(to->inbasket, result);
+ if(PyList_Append(to->inbasket, result) == -1){
+ Py_DECREF(result);
+ return NULL;
+ }
to->num_seen++;
lz->num_seen++;
return result;
@@ -66,18 +70,37 @@
to->outbasket = to->inbasket;
to->inbasket = tmp;
PyList_Reverse(to->outbasket);
- assert(PyList_GET_SIZE(to->outbasket) > 0);
}
+ /* Pop result from to->outbasket */
+ n = PyList_GET_SIZE(to->outbasket);
+ assert(n>0);
+ result = PyList_GET_ITEM(to->outbasket, n-1);
+ Py_INCREF(result);
+ if (PyList_SetSlice(to->outbasket, n-1, n, NULL) == -1) {
+ Py_DECREF(result);
+ return NULL;
+ }
lz->num_seen++;
- return PyObject_CallMethod(to->outbasket, "pop", NULL);
+ return result;
}
static void
ii_dealloc(iiobject *ii)
{
+ teeobject *to = ii->tee;
+ int n;
+
PyObject_GC_UnTrack(ii);
ii->tee->save_mode = 0; /* Stop saving data */
+ if (ii->num_seen < to->num_seen) {
+ /* It is the trailing iterator being freed; this means
+ that the stored queue data is no longer needed */
+ n = PyList_GET_SIZE(to->inbasket);
+ PyList_SetSlice(to->inbasket, 0, n, NULL);
+ n = PyList_GET_SIZE(to->outbasket);
+ PyList_SetSlice(to->outbasket, 0, n, NULL);
+ }
Py_XDECREF(ii->tee);
PyObject_GC_Del(ii);
}