Make built-in zip() equal to itertools.izip().

I mea, *really* equal -- for now, the implementation just imports
itertools. :-)
The only other changes necessary were various unit tests that were
assuming zip() returns a real list.  No "real" code made this assumption.

6 files changed, 63 insertions, 140 deletions

diff --git a/Lib/test/test_builtin.py b/Lib/test/test_builtin.py
index d3a272b..d4605e1 100644
--- a/Lib/test/test_builtin.py
+++ b/Lib/test/test_builtin.py
@@ -1554,18 +1554,18 @@ class BuiltinTest(unittest.TestCase):
         a = (1, 2, 3)
         b = (4, 5, 6)
         t = [(1, 4), (2, 5), (3, 6)]
-        self.assertEqual(zip(a, b), t)
+        self.assertEqual(list(zip(a, b)), t)
         b = [4, 5, 6]
-        self.assertEqual(zip(a, b), t)
+        self.assertEqual(list(zip(a, b)), t)
         b = (4, 5, 6, 7)
-        self.assertEqual(zip(a, b), t)
+        self.assertEqual(list(zip(a, b)), t)
         class I:
             def __getitem__(self, i):
                 if i < 0 or i > 2: raise IndexError
                 return i + 4
-        self.assertEqual(zip(a, I()), t)
-        self.assertEqual(zip(), [])
-        self.assertEqual(zip(*[]), [])
+        self.assertEqual(list(zip(a, I())), t)
+        self.assertEqual(list(zip()), [])
+        self.assertEqual(list(zip(*[])), [])
         self.assertRaises(TypeError, zip, None)
         class G:
             pass
@@ -1581,7 +1581,7 @@ class BuiltinTest(unittest.TestCase):
                 else:
                     return i
         self.assertEqual(
-            zip(SequenceWithoutALength(), xrange(2**30)),
+            list(zip(SequenceWithoutALength(), xrange(2**30))),
             list(enumerate(range(5)))
         )
 
@@ -1591,7 +1591,7 @@ class BuiltinTest(unittest.TestCase):
                     raise ValueError
                 else:
                     return i
-        self.assertRaises(ValueError, zip, BadSeq(), BadSeq())
+        self.assertRaises(ValueError, list, zip(BadSeq(), BadSeq()))
 
 class TestSorted(unittest.TestCase):
 
diff --git a/Lib/test/test_enumerate.py b/Lib/test/test_enumerate.py
index 10ff8f4..4ee8680 100644
--- a/Lib/test/test_enumerate.py
+++ b/Lib/test/test_enumerate.py
@@ -122,7 +122,7 @@ class TestEmpty(EnumerateTestCase):
 class TestBig(EnumerateTestCase):
 
     seq = range(10,20000,2)
-    res = zip(range(20000), seq)
+    res = list(zip(range(20000), seq))
 
 class TestReversed(unittest.TestCase):
 
diff --git a/Lib/test/test_iter.py b/Lib/test/test_iter.py
index 81be7e1..0ac34c8 100644
--- a/Lib/test/test_iter.py
+++ b/Lib/test/test_iter.py
@@ -423,21 +423,21 @@ class TestCase(unittest.TestCase):
 
     # Test zip()'s use of iterators.
     def test_builtin_zip(self):
-        self.assertEqual(zip(), [])
-        self.assertEqual(zip(*[]), [])
-        self.assertEqual(zip(*[(1, 2), 'ab']), [(1, 'a'), (2, 'b')])
+        self.assertEqual(list(zip()), [])
+        self.assertEqual(list(zip(*[])), [])
+        self.assertEqual(list(zip(*[(1, 2), 'ab'])), [(1, 'a'), (2, 'b')])
 
         self.assertRaises(TypeError, zip, None)
         self.assertRaises(TypeError, zip, range(10), 42)
         self.assertRaises(TypeError, zip, range(10), zip)
 
-        self.assertEqual(zip(IteratingSequenceClass(3)),
+        self.assertEqual(list(zip(IteratingSequenceClass(3))),
                          [(0,), (1,), (2,)])
-        self.assertEqual(zip(SequenceClass(3)),
+        self.assertEqual(list(zip(SequenceClass(3))),
                          [(0,), (1,), (2,)])
 
         d = {"one": 1, "two": 2, "three": 3}
-        self.assertEqual(d.items(), zip(d, d.itervalues()))
+        self.assertEqual(d.items(), list(zip(d, d.itervalues())))
 
         # Generate all ints starting at constructor arg.
         class IntsFrom:
@@ -459,7 +459,7 @@ class TestCase(unittest.TestCase):
             f.close()
         f = open(TESTFN, "r")
         try:
-            self.assertEqual(zip(IntsFrom(0), f, IntsFrom(-100)),
+            self.assertEqual(list(zip(IntsFrom(0), f, IntsFrom(-100))),
                              [(0, "a\n", -100),
                               (1, "bbb\n", -99),
                               (2, "cc\n", -98)])
@@ -470,7 +470,7 @@ class TestCase(unittest.TestCase):
             except OSError:
                 pass
 
-        self.assertEqual(zip(xrange(5)), [(i,) for i in range(5)])
+        self.assertEqual(list(zip(xrange(5))), [(i,) for i in range(5)])
 
         # Classes that lie about their lengths.
         class NoGuessLen5:
@@ -487,16 +487,19 @@ class TestCase(unittest.TestCase):
             def __len__(self):
                 return 30
 
+        def lzip(*args):
+            return list(zip(*args))
+
         self.assertEqual(len(Guess3Len5()), 3)
         self.assertEqual(len(Guess30Len5()), 30)
-        self.assertEqual(zip(NoGuessLen5()), zip(range(5)))
-        self.assertEqual(zip(Guess3Len5()), zip(range(5)))
-        self.assertEqual(zip(Guess30Len5()), zip(range(5)))
+        self.assertEqual(lzip(NoGuessLen5()), lzip(range(5)))
+        self.assertEqual(lzip(Guess3Len5()), lzip(range(5)))
+        self.assertEqual(lzip(Guess30Len5()), lzip(range(5)))
 
         expected = [(i, i) for i in range(5)]
         for x in NoGuessLen5(), Guess3Len5(), Guess30Len5():
             for y in NoGuessLen5(), Guess3Len5(), Guess30Len5():
-                self.assertEqual(zip(x, y), expected)
+                self.assertEqual(lzip(x, y), expected)
 
     # This test case will be removed if we don't have Unicode
     def test_unicode_join_endcase(self):
@@ -861,7 +864,7 @@ class TestCase(unittest.TestCase):
         a = range(5)
         e = enumerate(a)
         b = iter(e)
-        self.assertEqual(list(b), zip(range(5), range(5)))
+        self.assertEqual(list(b), list(zip(range(5), range(5))))
         self.assertEqual(list(b), [])
 
 
diff --git a/Lib/test/test_itertools.py b/Lib/test/test_itertools.py
index 29abd4b..8f8fc73 100644
--- a/Lib/test/test_itertools.py
+++ b/Lib/test/test_itertools.py
@@ -6,6 +6,9 @@ import sys
 import operator
 import random
 
+def lzip(*args):
+    return list(zip(*args))
+
 def onearg(x):
     'Test function of one argument'
     return 2*x
@@ -47,9 +50,9 @@ class TestBasicOps(unittest.TestCase):
         self.assertRaises(TypeError, chain, 2, 3)
 
     def test_count(self):
-        self.assertEqual(zip('abc',count()), [('a', 0), ('b', 1), ('c', 2)])
-        self.assertEqual(zip('abc',count(3)), [('a', 3), ('b', 4), ('c', 5)])
-        self.assertEqual(take(2, zip('abc',count(3))), [('a', 3), ('b', 4)])
+        self.assertEqual(lzip('abc',count()), [('a', 0), ('b', 1), ('c', 2)])
+        self.assertEqual(lzip('abc',count(3)), [('a', 3), ('b', 4), ('c', 5)])
+        self.assertEqual(take(2, lzip('abc',count(3))), [('a', 3), ('b', 4)])
         self.assertRaises(TypeError, count, 2, 3)
         self.assertRaises(TypeError, count, 'a')
         c = count(sys.maxint-2)   # verify that rollover doesn't crash
@@ -176,27 +179,28 @@ class TestBasicOps(unittest.TestCase):
         self.assertRaises(TypeError, ifilterfalse(range(6), range(6)).next)
 
     def test_izip(self):
+        # XXX This is rather silly now that builtin zip() calls izip()...
         ans = [(x,y) for x, y in izip('abc',count())]
         self.assertEqual(ans, [('a', 0), ('b', 1), ('c', 2)])
-        self.assertEqual(list(izip('abc', range(6))), zip('abc', range(6)))
-        self.assertEqual(list(izip('abcdef', range(3))), zip('abcdef', range(3)))
-        self.assertEqual(take(3,izip('abcdef', count())), zip('abcdef', range(3)))
-        self.assertEqual(list(izip('abcdef')), zip('abcdef'))
-        self.assertEqual(list(izip()), zip())
+        self.assertEqual(list(izip('abc', range(6))), lzip('abc', range(6)))
+        self.assertEqual(list(izip('abcdef', range(3))), lzip('abcdef', range(3)))
+        self.assertEqual(take(3,izip('abcdef', count())), lzip('abcdef', range(3)))
+        self.assertEqual(list(izip('abcdef')), lzip('abcdef'))
+        self.assertEqual(list(izip()), lzip())
         self.assertRaises(TypeError, izip, 3)
         self.assertRaises(TypeError, izip, range(3), 3)
         # Check tuple re-use (implementation detail)
         self.assertEqual([tuple(list(pair)) for pair in izip('abc', 'def')],
-                         zip('abc', 'def'))
+                         lzip('abc', 'def'))
         self.assertEqual([pair for pair in izip('abc', 'def')],
-                         zip('abc', 'def'))
+                         lzip('abc', 'def'))
         ids = map(id, izip('abc', 'def'))
         self.assertEqual(min(ids), max(ids))
         ids = map(id, list(izip('abc', 'def')))
         self.assertEqual(len(dict.fromkeys(ids)), len(ids))
 
     def test_repeat(self):
-        self.assertEqual(zip(xrange(3),repeat('a')),
+        self.assertEqual(lzip(xrange(3),repeat('a')),
                          [(0, 'a'), (1, 'a'), (2, 'a')])
         self.assertEqual(list(repeat('a', 3)), ['a', 'a', 'a'])
         self.assertEqual(take(3, repeat('a')), ['a', 'a', 'a'])
@@ -320,7 +324,7 @@ class TestBasicOps(unittest.TestCase):
         self.assertEqual(list(b), [])
 
         a, b = tee(irange(n)) # test 100% interleaved
-        self.assertEqual(zip(a,b), zip(range(n),range(n)))
+        self.assertEqual(lzip(a,b), lzip(range(n), range(n)))
 
         a, b = tee(irange(n)) # test 0% interleaved
         self.assertEqual(list(a), range(n))
@@ -601,8 +605,8 @@ class TestVariousIteratorArgs(unittest.TestCase):
     def test_izip(self):
         for s in ("123", "", range(1000), ('do', 1.2), xrange(2000,2200,5)):
             for g in (G, I, Ig, S, L, R):
-                self.assertEqual(list(izip(g(s))), zip(g(s)))
-                self.assertEqual(list(izip(g(s), g(s))), zip(g(s), g(s)))
+                self.assertEqual(list(izip(g(s))), lzip(g(s)))
+                self.assertEqual(list(izip(g(s), g(s))), lzip(g(s), g(s)))
             self.assertRaises(TypeError, izip, X(s))
             self.assertRaises(TypeError, izip, N(s))
             self.assertRaises(ZeroDivisionError, list, izip(E(s)))
@@ -627,7 +631,7 @@ class TestVariousIteratorArgs(unittest.TestCase):
     def test_starmap(self):
         for s in (range(10), range(0), range(100), (7,11), xrange(20,50,5)):
             for g in (G, I, Ig, S, L, R):
-                ss = zip(s, s)
+                ss = lzip(s, s)
                 self.assertEqual(list(starmap(operator.pow, g(ss))), map(operator.pow, g(s), g(s)))
             self.assertRaises(TypeError, starmap, operator.pow, X(ss))
             self.assertRaises(TypeError, starmap, operator.pow, N(ss))
diff --git a/Lib/test/test_tuple.py b/Lib/test/test_tuple.py
index dddb03e..96eb6e7 100644
--- a/Lib/test/test_tuple.py
+++ b/Lib/test/test_tuple.py
@@ -61,7 +61,7 @@ class TupleTest(seq_tests.CommonTest):
         base = range(N)
         xp = [(i, j) for i in base for j in base]
         inps = base + [(i, j) for i in base for j in xp] + \
-                     [(i, j) for i in xp for j in base] + xp + zip(base)
+                     [(i, j) for i in xp for j in base] + xp + list(zip(base))
         collisions = len(inps) - len(set(map(hash, inps)))
         self.assert_(collisions <= 15)
 
diff --git a/Python/bltinmodule.c b/Python/bltinmodule.c
index 94420f8..6ca2a28 100644
--- a/Python/bltinmodule.c
+++ b/Python/bltinmodule.c
@@ -1855,116 +1855,32 @@ is a shortcut for issubclass(X, A) or issubclass(X, B) or ... (etc.).");
 static PyObject*
 builtin_zip(PyObject *self, PyObject *args)
 {
-	PyObject *ret;
-	const Py_ssize_t itemsize = PySequence_Length(args);
-	Py_ssize_t i;
-	PyObject *itlist;  /* tuple of iterators */
-	Py_ssize_t len;	   /* guess at result length */
-
-	if (itemsize == 0)
-		return PyList_New(0);
-
-	/* args must be a tuple */
-	assert(PyTuple_Check(args));
-
-	/* Guess at result length:  the shortest of the input lengths.
-	   If some argument refuses to say, we refuse to guess too, lest
-	   an argument like xrange(sys.maxint) lead us astray.*/
-	len = -1;	/* unknown */
-	for (i = 0; i < itemsize; ++i) {
-		PyObject *item = PyTuple_GET_ITEM(args, i);
-		Py_ssize_t thislen = _PyObject_LengthHint(item);
-		if (thislen < 0) {
-			if (!PyErr_ExceptionMatches(PyExc_TypeError)  &&
-			    !PyErr_ExceptionMatches(PyExc_AttributeError)) {
-				return NULL;
-			}
-			PyErr_Clear();
-			len = -1;
-			break;
-		}
-		else if (len < 0 || thislen < len)
-			len = thislen;
-	}
-
-	/* allocate result list */
-	if (len < 0)
-		len = 10;	/* arbitrary */
-	if ((ret = PyList_New(len)) == NULL)
-		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 #%zd must support iteration",
-				    i+1);
-			goto Fail_ret_itlist;
-		}
-		PyTuple_SET_ITEM(itlist, i, it);
-	}
-
-	/* build result into ret list */
-	for (i = 0; ; ++i) {
-		int j;
-		PyObject *next = PyTuple_New(itemsize);
-		if (!next)
-			goto Fail_ret_itlist;
-
-		for (j = 0; j < itemsize; j++) {
-			PyObject *it = PyTuple_GET_ITEM(itlist, j);
-			PyObject *item = PyIter_Next(it);
-			if (!item) {
-				if (PyErr_Occurred()) {
-					Py_DECREF(ret);
-					ret = NULL;
-				}
-				Py_DECREF(next);
-				Py_DECREF(itlist);
-				goto Done;
-			}
-			PyTuple_SET_ITEM(next, j, item);
-		}
+	PyObject *itertools = NULL, *izip = NULL, *result = NULL;
 
-		if (i < len)
-			PyList_SET_ITEM(ret, i, next);
-		else {
-			int status = PyList_Append(ret, next);
-			Py_DECREF(next);
-			++len;
-			if (status < 0)
-				goto Fail_ret_itlist;
-		}
-	}
+	itertools = PyImport_ImportModule("itertools");
+	if (itertools == NULL)
+		return NULL;
+	
+	izip = PyObject_GetAttrString(itertools, "izip");
+	if (izip == NULL)
+		goto done;
 
-Done:
-	if (ret != NULL && i < len) {
-		/* The list is too big. */
-		if (PyList_SetSlice(ret, i, len, NULL) < 0)
-			return NULL;
-	}
-	return ret;
+	result = PyObject_Call(izip, args, NULL);
 
-Fail_ret_itlist:
-	Py_DECREF(itlist);
-Fail_ret:
-	Py_DECREF(ret);
-	return NULL;
+  done:
+	Py_XDECREF(itertools);
+	Py_XDECREF(izip);
+	return result;
 }
 
 
 PyDoc_STRVAR(zip_doc,
-"zip(seq1 [, seq2 [...]]) -> [(seq1[0], seq2[0] ...), (...)]\n\
+"zip(it1 [, it2 [...]]) -> iter([(it1[0], it2[0] ...), ...])\n\
 \n\
-Return a list of tuples, where each tuple contains the i-th element\n\
-from each of the argument sequences.  The returned list is truncated\n\
-in length to the length of the shortest argument sequence.");
+Return an iterator yielding tuples, where each tuple contains the\n\
+corresponding element from each of the argument iterables.\n\
+The returned iterator ends when the shortest argument iterable is exhausted.\n\
+NOTE: This is implemented using itertools.izip().");
 
 
 static PyMethodDef builtin_methods[] = {