Generalize zip() to work with iterators.

tim-one · tim-one · commit 8572b4fedf7e · 2001-05-06T01:05:02.000Z
NEEDS DOC CHANGES.
More AttributeErrors transmuted into TypeErrors, in test_b2.py, and,
again, this strikes me as a good thing.
This checkin completes the iterator generalization work that obviously
needed to be done.  Can anyone think of others that should be changed?
diff --git a/Lib/test/test_b2.py b/Lib/test/test_b2.py
@@ -309,13 +309,13 @@ class G:
 exc = 0
 try:
     zip(a, G())
-except AttributeError:
+except TypeError:
     exc = 1
 except:
     e = sys.exc_info()[0]
     raise TestFailed, 'zip(a, b) - b instance w/o __getitem__'
 if not exc:
-    raise TestFailed, 'zip(a, b) - missing expected AttributeError'
+    raise TestFailed, 'zip(a, b) - missing expected TypeError'
 
 
 # Epilogue -- unlink the temp file
diff --git a/Lib/test/test_iter.py b/Lib/test/test_iter.py
@@ -418,6 +418,52 @@ def test_builtin_map(self):
             except OSError:
                 pass
 
+    # Test zip()'s use of iterators.
+    def test_builtin_zip(self):
+        self.assertRaises(TypeError, zip)
+        self.assertRaises(TypeError, zip, None)
+        self.assertRaises(TypeError, zip, range(10), 42)
+        self.assertRaises(TypeError, zip, range(10), zip)
+
+        self.assertEqual(zip(IteratingSequenceClass(3)),
+                         [(0,), (1,), (2,)])
+        self.assertEqual(zip(SequenceClass(3)),
+                         [(0,), (1,), (2,)])
+
+        d = {"one": 1, "two": 2, "three": 3}
+        self.assertEqual(d.items(), zip(d, d.itervalues()))
+
+        # Generate all ints starting at constructor arg.
+        class IntsFrom:
+            def __init__(self, start):
+                self.i = start
+
+            def __iter__(self):
+                return self
+
+            def next(self):
+                i = self.i
+                self.i = i+1
+                return i
+
+        f = open(TESTFN, "w")
+        try:
+            f.write("a\n" "bbb\n" "cc\n")
+        finally:
+            f.close()
+        f = open(TESTFN, "r")
+        try:
+            self.assertEqual(zip(IntsFrom(0), f, IntsFrom(-100)),
+                             [(0, "a\n", -100),
+                              (1, "bbb\n", -99),
+                              (2, "cc\n", -98)])
+        finally:
+            f.close()
+            try:
+                unlink(TESTFN)
+            except OSError:
+                pass
+
     # Test reduces()'s use of iterators.
     def test_builtin_reduce(self):
         from operator import add
diff --git a/Misc/NEWS b/Misc/NEWS
@@ -17,17 +17,13 @@ Core
 
 - The following functions were generalized to work nicely with iterator
   arguments:
-    filter()
-    list()
-    map()
-    max()
-    min()
-    reduce()
-    tuple() (PySequence_Tuple() and PySequence_Fast() in C API)
+    map(), filter(), reduce()
+    list(), tuple() (PySequence_Tuple() and PySequence_Fast() in C API)
+    max(), min()
+    zip()
     .join() method of strings
     'x in y' and 'x not in y' (PySequence_Contains() in C API)
     operator.countOf() (PySequence_Count() in C API)
-    XXX TODO zip()
 
 
 What's New in Python 2.1 (final)?
diff --git a/Python/bltinmodule.c b/Python/bltinmodule.c
@@ -2102,7 +2102,8 @@ builtin_zip(PyObject *self, PyObject *args)
 {
 	PyObject *ret;
 	int itemsize = PySequence_Length(args);
-	int i, j;
+	int i;
+	PyObject *itlist;  /* tuple of iterators */
 
 	if (itemsize < 1) {
 		PyErr_SetString(PyExc_TypeError,
@@ -2112,35 +2113,60 @@ builtin_zip(PyObject *self, PyObject *args)
 	/* args must be a tuple */
 	assert(PyTuple_Check(args));
 
+	/* allocate result list */
 	if ((ret = PyList_New(0)) == NULL)
 		return NULL;
 
-	for (i = 0;; i++) {
-		PyObject *next = PyTuple_New(itemsize);
-		if (!next) {
-			Py_DECREF(ret);
-			return NULL;
+	/* obtain iterators */
+	itlist = PyTuple_New(itemsize);
+	if (itlist == NULL)
+		goto Fail_ret;
+	for (i = 0; i < itemsize; ++i) {
+		PyObject *item = PyTuple_GET_ITEM(args, i);
+		PyObject *it = PyObject_GetIter(item);
+		if (it == NULL) {
+			if (PyErr_ExceptionMatches(PyExc_TypeError))
+				PyErr_Format(PyExc_TypeError,
+				    "zip argument #%d must support iteration",
+				    i+1);
+			goto Fail_ret_itlist;
 		}
-		for (j = 0; j < itemsize; j++) {
-			PyObject *seq = PyTuple_GET_ITEM(args, j);
-			PyObject *item = PySequence_GetItem(seq, i);
+		PyTuple_SET_ITEM(itlist, i, it);
+	}
 
+	/* build result into ret list */
+	for (;;) {
+		int status;
+		PyObject *next = PyTuple_New(itemsize);
+		if (!next)
+			goto Fail_ret_itlist;
+
+		for (i = 0; i < itemsize; i++) {
+			PyObject *it = PyTuple_GET_ITEM(itlist, i);
+			PyObject *item = PyIter_Next(it);
 			if (!item) {
-				if (PyErr_ExceptionMatches(PyExc_IndexError)) {
-					PyErr_Clear();
-					Py_DECREF(next);
-					return ret;
+				if (PyErr_Occurred()) {
+					Py_DECREF(ret);
+					ret = NULL;
 				}
 				Py_DECREF(next);
-				Py_DECREF(ret);
-				return NULL;
+				Py_DECREF(itlist);
+				return ret;
 			}
-			PyTuple_SET_ITEM(next, j, item);
+			PyTuple_SET_ITEM(next, i, item);
 		}
-		PyList_Append(ret, next);
+
+		status = PyList_Append(ret, next);
 		Py_DECREF(next);
+		if (status < 0)
+			goto Fail_ret_itlist;
 	}
-	/* no return */
+
+Fail_ret_itlist:
+	Py_DECREF(itlist);
+Fail_ret:
+	Py_DECREF(ret);
+	return NULL;
 }
 
 
