[Python-checkins] cpython: Issue #24874: Speed-up itertools and make it pickles more compact.

Sun Aug 16 23:49:38 CEST 2015

https://hg.python.org/cpython/rev/17b5c8ba6875
changeset: 97400:17b5c8ba6875
user: Raymond Hettinger <python at rcn.com>
date: Sun Aug 16 14:49:24 2015 -0700
summary:
 Issue #24874: Speed-up itertools and make it pickles more compact.
files:
 Misc/NEWS | 3 +
 Modules/itertoolsmodule.c | 64 +++++++++++++++++---------
 2 files changed, 45 insertions(+), 22 deletions(-)

diff --git a/Misc/NEWS b/Misc/NEWS
--- a/Misc/NEWS
+++ b/Misc/NEWS
@@ -25,6 +25,9 @@
 - Issue #21159: Improve message in configparser.InterpolationMissingOptionError.
 Patch from Łukasz Langa.
 
+- Issue #24874: Improve speed of itertools.cycle() and make its
+ pickle more compact.
+
 - Fix crash in itertools.cycle.__setstate__() when the first argument wasn't
 a list.
 
diff --git a/Modules/itertoolsmodule.c b/Modules/itertoolsmodule.c
--- a/Modules/itertoolsmodule.c
+++ b/Modules/itertoolsmodule.c
@@ -1,3 +1,5 @@
+
+#define PY_SSIZE_T_CLEAN
 
 #include "Python.h"
 #include "structmember.h"
@@ -863,6 +865,7 @@
 PyObject_HEAD
 PyObject *it;
 PyObject *saved;
+ Py_ssize_t index;
 int firstpass;
 } cycleobject;
 
@@ -902,6 +905,7 @@
 }
 lz->it = it;
 lz->saved = saved;
+ lz->index = 0;
 lz->firstpass = 0;
 
 return (PyObject *)lz;
@@ -911,15 +915,16 @@
 cycle_dealloc(cycleobject *lz)
 {
 PyObject_GC_UnTrack(lz);
+ Py_XDECREF(lz->it);
 Py_XDECREF(lz->saved);
- Py_XDECREF(lz->it);
 Py_TYPE(lz)->tp_free(lz);
 }
 
 static int
 cycle_traverse(cycleobject *lz, visitproc visit, void *arg)
 {
- Py_VISIT(lz->it);
+ if (lz->it)
+ Py_VISIT(lz->it);
 Py_VISIT(lz->saved);
 return 0;
 }
@@ -928,13 +933,13 @@
 cycle_next(cycleobject *lz)
 {
 PyObject *item;
- PyObject *it;
- PyObject *tmp;
-
- while (1) {
+
+ if (lz->it != NULL) {
 item = PyIter_Next(lz->it);
 if (item != NULL) {
- if (!lz->firstpass && PyList_Append(lz->saved, item)) {
+ if (lz->firstpass)
+ return item;
+ if (PyList_Append(lz->saved, item)) {
 Py_DECREF(item);
 return NULL;
 }
@@ -942,27 +947,41 @@
 }
 /* Note: StopIteration is already cleared by PyIter_Next() */
 if (PyErr_Occurred())
- return NULL;
- if (PyList_Size(lz->saved) == 0)
 return NULL;
- it = PyObject_GetIter(lz->saved);
- if (it == NULL)
- return NULL;
- tmp = lz->it;
- lz->it = it;
- lz->firstpass = 1;
- Py_DECREF(tmp);
+ Py_CLEAR(lz->it);
 }
+ if (Py_SIZE(lz->saved) == 0)
+ return NULL;
+ item = PyList_GET_ITEM(lz->saved, lz->index);
+ lz->index++;
+ if (lz->index >= Py_SIZE(lz->saved))
+ lz->index = 0;
+ Py_INCREF(item);
+ return item;
 }
 
 static PyObject *
 cycle_reduce(cycleobject *lz)
 {
- /* Create a new cycle with the iterator tuple, then set
- * the saved state on it.
- */
- return Py_BuildValue("O(O)(Oi)", Py_TYPE(lz),
- lz->it, lz->saved, lz->firstpass);
+ /* Create a new cycle with the iterator tuple, then set the saved state */
+ if (lz->it == NULL) {
+ PyObject *it = PyObject_GetIter(lz->saved);
+ if (it == NULL)
+ return NULL;
+ if (lz->index != 0) {
+ _Py_IDENTIFIER(__setstate__);
+ PyObject *res = _PyObject_CallMethodId(it, &PyId___setstate__,
+ "n", lz->index);
+ if (res == NULL) {
+ Py_DECREF(it);
+ return NULL;
+ }
+ Py_DECREF(res);
+ }
+ return Py_BuildValue("O(N)(Oi)", Py_TYPE(lz), it, lz->saved, 1);
+ }
+ return Py_BuildValue("O(O)(Oi)", Py_TYPE(lz), lz->it, lz->saved,
+ lz->firstpass);
 }
 
 static PyObject *
@@ -972,10 +991,11 @@
 int firstpass;
 if (!PyArg_ParseTuple(state, "O!i", &PyList_Type, &saved, &firstpass))
 return NULL;
+ Py_INCREF(saved);
 Py_CLEAR(lz->saved);
 lz->saved = saved;
- Py_XINCREF(lz->saved);
 lz->firstpass = firstpass != 0;
+ lz->index = 0;
 Py_RETURN_NONE;
 }
 
-- 
Repository URL: https://hg.python.org/cpython
