Generalize zip() to work with iterators.

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
index 28022153b560818fcbbf10b10775255cbae3c55a..324d02fdd1a208e0d484222540eba53f971fa0ba 100644 (file)
--- 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
index a50c74f2fd58c183f8a57bb82defcf28c987d643..ddc58a7cfc727a4aa2cadc74c512de4511ebdaa6 100644 (file)
--- a/Lib/test/test_iter.py
+++ b/Lib/test/test_iter.py
@@ -418,6 +418,52 @@ class TestCase(unittest.TestCase):
             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
index aecc5e9114c4254996177265ba97de59cf0094b8..1f971cd40c94554dad02d2d604e59c161ce8bf6e 100644 (file)
--- 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
index 4a51ccd98a1279287fe3f4511ae3090370033eeb..cc1bc957a96cae7735a26d6cdc09147c8e7bcc14 100644 (file)
--- 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;
 }