SF patch 1546297 (with some tweaks):

Create a real zip iterator object; not using itertools.izip
(Brian Holmes).

1 files changed, 171 insertions, 0 deletions

diff --git a/Objects/iterobject.c b/Objects/iterobject.c
index cf839f4..dc4c583 100644
--- a/Objects/iterobject.c
+++ b/Objects/iterobject.c
@@ -230,3 +230,174 @@ PyTypeObject PyCallIter_Type = {
 	(iternextfunc)calliter_iternext,	/* tp_iternext */
 	0,					/* tp_methods */
 };
+
+
+/*********************** Zip Iterator **************************/
+/* Largely copied from itertools.c by Brian Holmes */
+
+typedef struct zipiterobject_t {
+	PyObject_HEAD
+	PyTupleObject *it_tuple;  /* Set to NULL when iterator is exhausted */
+        Py_ssize_t resultsize;	
+	PyTupleObject *result;	/* Reusable tuple for optimization */
+} zipiterobject;
+
+static PyTypeObject PyZipIter_Type; /* Forward */
+
+PyObject *
+_PyZip_CreateIter(PyObject* args)
+{
+	Py_ssize_t i;
+        Py_ssize_t tuplesize;
+	PyObject* ziptuple;
+	PyObject* result;
+	struct zipiterobject_t* zipiter;
+        
+        assert(PyTuple_Check(args));
+
+	if (PyZipIter_Type.ob_type == NULL) {
+		if (PyType_Ready(&PyZipIter_Type) < 0)
+			return NULL;
+	}
+
+	tuplesize = PySequence_Length((PyObject*) args);
+
+	ziptuple = PyTuple_New(tuplesize);
+	if (ziptuple == NULL)
+		return NULL;
+
+	for (i = 0; i < tuplesize; i++) {
+		PyObject *o = PyTuple_GET_ITEM(args, i);
+		PyObject *it = PyObject_GetIter(o);
+		if (it == NULL) {
+			/* XXX Should we do this?
+			if (PyErr_ExceptionMatches(PyExc_TypeError))
+				PyErr_Format(PyExc_TypeError, 
+				  "zip argument #%zd must support iteration",
+					I+1);
+			*/
+			Py_DECREF(ziptuple);
+			return NULL;
+		}
+		PyTuple_SET_ITEM(ziptuple, i, it);
+	}
+
+        /* create a reusable result holder */
+        result = PyTuple_New(tuplesize);
+        if (result == NULL) {
+                Py_DECREF(ziptuple);
+                return NULL;
+        }
+        for (i = 0; i < tuplesize; i++) {
+                Py_INCREF(Py_None);
+                PyTuple_SET_ITEM(result, i, Py_None);
+        }
+	
+	zipiter = PyObject_GC_New(zipiterobject, &PyZipIter_Type);
+	if (zipiter == NULL) {
+		Py_DECREF(ziptuple);
+		Py_DECREF(result);
+		return NULL;
+	}
+
+	zipiter->result = (PyTupleObject*) result;
+        zipiter->resultsize = tuplesize;
+	Py_INCREF(ziptuple);
+	zipiter->it_tuple = (PyTupleObject *) ziptuple;
+	_PyObject_GC_TRACK(zipiter);
+	return (PyObject *)zipiter;
+}
+
+static void
+zipiter_dealloc(zipiterobject *it)
+{
+	_PyObject_GC_UNTRACK(it);
+	Py_XDECREF(it->it_tuple);
+	Py_XDECREF(it->result);
+	PyObject_GC_Del(it);
+}
+
+static int
+zipiter_traverse(zipiterobject *it, visitproc visit, void *arg)
+{
+	Py_VISIT(it->it_tuple);
+	Py_VISIT(it->result);
+	return 0;
+}
+
+static PyObject *
+zipiter_next(zipiterobject *zit)
+{
+        Py_ssize_t i;
+        Py_ssize_t tuplesize = zit->resultsize;
+        PyObject *result = (PyObject*) zit->result;
+        PyObject *olditem;
+
+        if (tuplesize == 0)
+                return NULL;
+
+        if (result->ob_refcnt == 1) {
+		Py_INCREF(result);
+		for (i = 0; i < tuplesize; i++) {
+			PyObject *it = PyTuple_GET_ITEM(zit->it_tuple, i);
+			assert(PyIter_Check(it));
+			PyObject *item = (*it->ob_type->tp_iternext)(it);
+			if (item == NULL) {
+				Py_DECREF(result);
+				return NULL;
+			}
+			olditem = PyTuple_GET_ITEM(result, i);
+			PyTuple_SET_ITEM(result, i, item);
+			Py_DECREF(olditem);
+		}
+	} else {
+		result = PyTuple_New(tuplesize);
+		if (result == NULL)
+			return NULL;
+		for (i = 0; i < tuplesize; i++) {
+			PyObject *it = PyTuple_GET_ITEM(zit->it_tuple, i);
+			assert(PyIter_Check(it));
+			PyObject *item = (*it->ob_type->tp_iternext)(it);
+			if (item == NULL) {
+				Py_DECREF(result);
+				return NULL;
+			}
+			PyTuple_SET_ITEM(result, i, item);
+		}
+	}
+	return result;
+}
+
+static PyTypeObject PyZipIter_Type = {
+	PyObject_HEAD_INIT(0)
+	0,					/* ob_size */
+	"zipiterator",				/* tp_name */
+	sizeof(zipiterobject),			/* tp_basicsize */
+	0,					/* tp_itemsize */
+	/* methods */
+	(destructor)zipiter_dealloc,		/* tp_dealloc */
+	0,					/* tp_print */
+	0,					/* tp_getattr */
+	0,					/* tp_setattr */
+	0,					/* tp_compare */
+	0,					/* tp_repr */
+	0,					/* tp_as_number */
+	0,					/* tp_as_sequence */
+	0,					/* tp_as_mapping */
+	0,					/* tp_hash */
+	0,					/* tp_call */
+	0,					/* tp_str */
+	PyObject_GenericGetAttr,		/* tp_getattro */
+	0,					/* tp_setattro */
+	0,					/* tp_as_buffer */
+	Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC,/* tp_flags */
+	0,					/* tp_doc */
+	(traverseproc)zipiter_traverse,		/* tp_traverse */
+	0,					/* tp_clear */
+	0,					/* tp_richcompare */
+	0,					/* tp_weakzipoffset */
+	PyObject_SelfIter,			/* tp_iter */
+	(iternextfunc)zipiter_next,		/* tp_iternext */
+	0,					/* tp_methods */
+	0,					/* tp_members */
+};