SF patch 1546297 (with some tweaks):

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

3 files changed, 177 insertions, 15 deletions

diff --git a/Include/iterobject.h b/Include/iterobject.h
index c078ebb..69deb45 100644
--- a/Include/iterobject.h
+++ b/Include/iterobject.h
@@ -16,6 +16,9 @@ PyAPI_DATA(PyTypeObject) PyCallIter_Type;
 #define PyCallIter_Check(op) ((op)->ob_type == &PyCallIter_Type)
 
 PyAPI_FUNC(PyObject *) PyCallIter_New(PyObject *, PyObject *);
+
+PyObject* _PyZip_CreateIter(PyObject* args);
+
 #ifdef __cplusplus
 }
 #endif
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 */
+};
diff --git a/Python/bltinmodule.c b/Python/bltinmodule.c
index 6ca2a28..200ec26 100644
--- a/Python/bltinmodule.c
+++ b/Python/bltinmodule.c
@@ -1855,22 +1855,10 @@ is a shortcut for issubclass(X, A) or issubclass(X, B) or ... (etc.).");
 static PyObject*
 builtin_zip(PyObject *self, PyObject *args)
 {
-	PyObject *itertools = NULL, *izip = NULL, *result = NULL;
+	/* args must be a tuple */
+	assert(PyTuple_Check(args));
 
-	itertools = PyImport_ImportModule("itertools");
-	if (itertools == NULL)
-		return NULL;
-	
-	izip = PyObject_GetAttrString(itertools, "izip");
-	if (izip == NULL)
-		goto done;
-
-	result = PyObject_Call(izip, args, NULL);
-
-  done:
-	Py_XDECREF(itertools);
-	Py_XDECREF(izip);
-	return result;
+	return _PyZip_CreateIter(args);
 }