Commit 801f0d78 authored by Guido van Rossum's avatar Guido van Rossum

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.
parent d38abe94
...@@ -1554,18 +1554,18 @@ class BuiltinTest(unittest.TestCase): ...@@ -1554,18 +1554,18 @@ class BuiltinTest(unittest.TestCase):
a = (1, 2, 3) a = (1, 2, 3)
b = (4, 5, 6) b = (4, 5, 6)
t = [(1, 4), (2, 5), (3, 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] b = [4, 5, 6]
self.assertEqual(zip(a, b), t) self.assertEqual(list(zip(a, b)), t)
b = (4, 5, 6, 7) b = (4, 5, 6, 7)
self.assertEqual(zip(a, b), t) self.assertEqual(list(zip(a, b)), t)
class I: class I:
def __getitem__(self, i): def __getitem__(self, i):
if i < 0 or i > 2: raise IndexError if i < 0 or i > 2: raise IndexError
return i + 4 return i + 4
self.assertEqual(zip(a, I()), t) self.assertEqual(list(zip(a, I())), t)
self.assertEqual(zip(), []) self.assertEqual(list(zip()), [])
self.assertEqual(zip(*[]), []) self.assertEqual(list(zip(*[])), [])
self.assertRaises(TypeError, zip, None) self.assertRaises(TypeError, zip, None)
class G: class G:
pass pass
...@@ -1581,7 +1581,7 @@ class BuiltinTest(unittest.TestCase): ...@@ -1581,7 +1581,7 @@ class BuiltinTest(unittest.TestCase):
else: else:
return i return i
self.assertEqual( self.assertEqual(
zip(SequenceWithoutALength(), xrange(2**30)), list(zip(SequenceWithoutALength(), xrange(2**30))),
list(enumerate(range(5))) list(enumerate(range(5)))
) )
...@@ -1591,7 +1591,7 @@ class BuiltinTest(unittest.TestCase): ...@@ -1591,7 +1591,7 @@ class BuiltinTest(unittest.TestCase):
raise ValueError raise ValueError
else: else:
return i return i
self.assertRaises(ValueError, zip, BadSeq(), BadSeq()) self.assertRaises(ValueError, list, zip(BadSeq(), BadSeq()))
class TestSorted(unittest.TestCase): class TestSorted(unittest.TestCase):
......
...@@ -122,7 +122,7 @@ class TestEmpty(EnumerateTestCase): ...@@ -122,7 +122,7 @@ class TestEmpty(EnumerateTestCase):
class TestBig(EnumerateTestCase): class TestBig(EnumerateTestCase):
seq = range(10,20000,2) seq = range(10,20000,2)
res = zip(range(20000), seq) res = list(zip(range(20000), seq))
class TestReversed(unittest.TestCase): class TestReversed(unittest.TestCase):
......
...@@ -423,21 +423,21 @@ class TestCase(unittest.TestCase): ...@@ -423,21 +423,21 @@ class TestCase(unittest.TestCase):
# Test zip()'s use of iterators. # Test zip()'s use of iterators.
def test_builtin_zip(self): def test_builtin_zip(self):
self.assertEqual(zip(), []) self.assertEqual(list(zip()), [])
self.assertEqual(zip(*[]), []) self.assertEqual(list(zip(*[])), [])
self.assertEqual(zip(*[(1, 2), 'ab']), [(1, 'a'), (2, 'b')]) self.assertEqual(list(zip(*[(1, 2), 'ab'])), [(1, 'a'), (2, 'b')])
self.assertRaises(TypeError, zip, None) self.assertRaises(TypeError, zip, None)
self.assertRaises(TypeError, zip, range(10), 42) self.assertRaises(TypeError, zip, range(10), 42)
self.assertRaises(TypeError, zip, range(10), zip) self.assertRaises(TypeError, zip, range(10), zip)
self.assertEqual(zip(IteratingSequenceClass(3)), self.assertEqual(list(zip(IteratingSequenceClass(3))),
[(0,), (1,), (2,)]) [(0,), (1,), (2,)])
self.assertEqual(zip(SequenceClass(3)), self.assertEqual(list(zip(SequenceClass(3))),
[(0,), (1,), (2,)]) [(0,), (1,), (2,)])
d = {"one": 1, "two": 2, "three": 3} 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. # Generate all ints starting at constructor arg.
class IntsFrom: class IntsFrom:
...@@ -459,7 +459,7 @@ class TestCase(unittest.TestCase): ...@@ -459,7 +459,7 @@ class TestCase(unittest.TestCase):
f.close() f.close()
f = open(TESTFN, "r") f = open(TESTFN, "r")
try: try:
self.assertEqual(zip(IntsFrom(0), f, IntsFrom(-100)), self.assertEqual(list(zip(IntsFrom(0), f, IntsFrom(-100))),
[(0, "a\n", -100), [(0, "a\n", -100),
(1, "bbb\n", -99), (1, "bbb\n", -99),
(2, "cc\n", -98)]) (2, "cc\n", -98)])
...@@ -470,7 +470,7 @@ class TestCase(unittest.TestCase): ...@@ -470,7 +470,7 @@ class TestCase(unittest.TestCase):
except OSError: except OSError:
pass 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. # Classes that lie about their lengths.
class NoGuessLen5: class NoGuessLen5:
...@@ -487,16 +487,19 @@ class TestCase(unittest.TestCase): ...@@ -487,16 +487,19 @@ class TestCase(unittest.TestCase):
def __len__(self): def __len__(self):
return 30 return 30
def lzip(*args):
return list(zip(*args))
self.assertEqual(len(Guess3Len5()), 3) self.assertEqual(len(Guess3Len5()), 3)
self.assertEqual(len(Guess30Len5()), 30) self.assertEqual(len(Guess30Len5()), 30)
self.assertEqual(zip(NoGuessLen5()), zip(range(5))) self.assertEqual(lzip(NoGuessLen5()), lzip(range(5)))
self.assertEqual(zip(Guess3Len5()), zip(range(5))) self.assertEqual(lzip(Guess3Len5()), lzip(range(5)))
self.assertEqual(zip(Guess30Len5()), zip(range(5))) self.assertEqual(lzip(Guess30Len5()), lzip(range(5)))
expected = [(i, i) for i in range(5)] expected = [(i, i) for i in range(5)]
for x in NoGuessLen5(), Guess3Len5(), Guess30Len5(): for x in NoGuessLen5(), Guess3Len5(), Guess30Len5():
for y 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 # This test case will be removed if we don't have Unicode
def test_unicode_join_endcase(self): def test_unicode_join_endcase(self):
...@@ -861,7 +864,7 @@ class TestCase(unittest.TestCase): ...@@ -861,7 +864,7 @@ class TestCase(unittest.TestCase):
a = range(5) a = range(5)
e = enumerate(a) e = enumerate(a)
b = iter(e) 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), []) self.assertEqual(list(b), [])
......
...@@ -6,6 +6,9 @@ import sys ...@@ -6,6 +6,9 @@ import sys
import operator import operator
import random import random
def lzip(*args):
return list(zip(*args))
def onearg(x): def onearg(x):
'Test function of one argument' 'Test function of one argument'
return 2*x return 2*x
...@@ -47,9 +50,9 @@ class TestBasicOps(unittest.TestCase): ...@@ -47,9 +50,9 @@ class TestBasicOps(unittest.TestCase):
self.assertRaises(TypeError, chain, 2, 3) self.assertRaises(TypeError, chain, 2, 3)
def test_count(self): def test_count(self):
self.assertEqual(zip('abc',count()), [('a', 0), ('b', 1), ('c', 2)]) self.assertEqual(lzip('abc',count()), [('a', 0), ('b', 1), ('c', 2)])
self.assertEqual(zip('abc',count(3)), [('a', 3), ('b', 4), ('c', 5)]) self.assertEqual(lzip('abc',count(3)), [('a', 3), ('b', 4), ('c', 5)])
self.assertEqual(take(2, zip('abc',count(3))), [('a', 3), ('b', 4)]) self.assertEqual(take(2, lzip('abc',count(3))), [('a', 3), ('b', 4)])
self.assertRaises(TypeError, count, 2, 3) self.assertRaises(TypeError, count, 2, 3)
self.assertRaises(TypeError, count, 'a') self.assertRaises(TypeError, count, 'a')
c = count(sys.maxint-2) # verify that rollover doesn't crash c = count(sys.maxint-2) # verify that rollover doesn't crash
...@@ -176,27 +179,28 @@ class TestBasicOps(unittest.TestCase): ...@@ -176,27 +179,28 @@ class TestBasicOps(unittest.TestCase):
self.assertRaises(TypeError, ifilterfalse(range(6), range(6)).next) self.assertRaises(TypeError, ifilterfalse(range(6), range(6)).next)
def test_izip(self): 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())] ans = [(x,y) for x, y in izip('abc',count())]
self.assertEqual(ans, [('a', 0), ('b', 1), ('c', 2)]) self.assertEqual(ans, [('a', 0), ('b', 1), ('c', 2)])
self.assertEqual(list(izip('abc', range(6))), zip('abc', range(6))) self.assertEqual(list(izip('abc', range(6))), lzip('abc', range(6)))
self.assertEqual(list(izip('abcdef', range(3))), zip('abcdef', range(3))) self.assertEqual(list(izip('abcdef', range(3))), lzip('abcdef', range(3)))
self.assertEqual(take(3,izip('abcdef', count())), zip('abcdef', range(3))) self.assertEqual(take(3,izip('abcdef', count())), lzip('abcdef', range(3)))
self.assertEqual(list(izip('abcdef')), zip('abcdef')) self.assertEqual(list(izip('abcdef')), lzip('abcdef'))
self.assertEqual(list(izip()), zip()) self.assertEqual(list(izip()), lzip())
self.assertRaises(TypeError, izip, 3) self.assertRaises(TypeError, izip, 3)
self.assertRaises(TypeError, izip, range(3), 3) self.assertRaises(TypeError, izip, range(3), 3)
# Check tuple re-use (implementation detail) # Check tuple re-use (implementation detail)
self.assertEqual([tuple(list(pair)) for pair in izip('abc', 'def')], 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')], self.assertEqual([pair for pair in izip('abc', 'def')],
zip('abc', 'def')) lzip('abc', 'def'))
ids = map(id, izip('abc', 'def')) ids = map(id, izip('abc', 'def'))
self.assertEqual(min(ids), max(ids)) self.assertEqual(min(ids), max(ids))
ids = map(id, list(izip('abc', 'def'))) ids = map(id, list(izip('abc', 'def')))
self.assertEqual(len(dict.fromkeys(ids)), len(ids)) self.assertEqual(len(dict.fromkeys(ids)), len(ids))
def test_repeat(self): def test_repeat(self):
self.assertEqual(zip(xrange(3),repeat('a')), self.assertEqual(lzip(xrange(3),repeat('a')),
[(0, 'a'), (1, 'a'), (2, 'a')]) [(0, 'a'), (1, 'a'), (2, 'a')])
self.assertEqual(list(repeat('a', 3)), ['a', 'a', 'a']) self.assertEqual(list(repeat('a', 3)), ['a', 'a', 'a'])
self.assertEqual(take(3, repeat('a')), ['a', 'a', 'a']) self.assertEqual(take(3, repeat('a')), ['a', 'a', 'a'])
...@@ -320,7 +324,7 @@ class TestBasicOps(unittest.TestCase): ...@@ -320,7 +324,7 @@ class TestBasicOps(unittest.TestCase):
self.assertEqual(list(b), []) self.assertEqual(list(b), [])
a, b = tee(irange(n)) # test 100% interleaved 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 a, b = tee(irange(n)) # test 0% interleaved
self.assertEqual(list(a), range(n)) self.assertEqual(list(a), range(n))
...@@ -601,8 +605,8 @@ class TestVariousIteratorArgs(unittest.TestCase): ...@@ -601,8 +605,8 @@ class TestVariousIteratorArgs(unittest.TestCase):
def test_izip(self): def test_izip(self):
for s in ("123", "", range(1000), ('do', 1.2), xrange(2000,2200,5)): for s in ("123", "", range(1000), ('do', 1.2), xrange(2000,2200,5)):
for g in (G, I, Ig, S, L, R): for g in (G, I, Ig, S, L, R):
self.assertEqual(list(izip(g(s))), zip(g(s))) self.assertEqual(list(izip(g(s))), lzip(g(s)))
self.assertEqual(list(izip(g(s), g(s))), zip(g(s), 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, X(s))
self.assertRaises(TypeError, izip, N(s)) self.assertRaises(TypeError, izip, N(s))
self.assertRaises(ZeroDivisionError, list, izip(E(s))) self.assertRaises(ZeroDivisionError, list, izip(E(s)))
...@@ -627,7 +631,7 @@ class TestVariousIteratorArgs(unittest.TestCase): ...@@ -627,7 +631,7 @@ class TestVariousIteratorArgs(unittest.TestCase):
def test_starmap(self): def test_starmap(self):
for s in (range(10), range(0), range(100), (7,11), xrange(20,50,5)): for s in (range(10), range(0), range(100), (7,11), xrange(20,50,5)):
for g in (G, I, Ig, S, L, R): 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.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, X(ss))
self.assertRaises(TypeError, starmap, operator.pow, N(ss)) self.assertRaises(TypeError, starmap, operator.pow, N(ss))
......
...@@ -61,7 +61,7 @@ class TupleTest(seq_tests.CommonTest): ...@@ -61,7 +61,7 @@ class TupleTest(seq_tests.CommonTest):
base = range(N) base = range(N)
xp = [(i, j) for i in base for j in base] xp = [(i, j) for i in base for j in base]
inps = base + [(i, j) for i in base for j in xp] + \ 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))) collisions = len(inps) - len(set(map(hash, inps)))
self.assert_(collisions <= 15) self.assert_(collisions <= 15)
......
...@@ -1855,116 +1855,32 @@ is a shortcut for issubclass(X, A) or issubclass(X, B) or ... (etc.)."); ...@@ -1855,116 +1855,32 @@ is a shortcut for issubclass(X, A) or issubclass(X, B) or ... (etc.).");
static PyObject* static PyObject*
builtin_zip(PyObject *self, PyObject *args) builtin_zip(PyObject *self, PyObject *args)
{ {
PyObject *ret; PyObject *itertools = NULL, *izip = NULL, *result = NULL;
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);
}
if (i < len) itertools = PyImport_ImportModule("itertools");
PyList_SET_ITEM(ret, i, next); if (itertools == NULL)
else { return NULL;
int status = PyList_Append(ret, next);
Py_DECREF(next); izip = PyObject_GetAttrString(itertools, "izip");
++len; if (izip == NULL)
if (status < 0) goto done;
goto Fail_ret_itlist;
}
}
Done: result = PyObject_Call(izip, args, NULL);
if (ret != NULL && i < len) {
/* The list is too big. */
if (PyList_SetSlice(ret, i, len, NULL) < 0)
return NULL;
}
return ret;
Fail_ret_itlist: done:
Py_DECREF(itlist); Py_XDECREF(itertools);
Fail_ret: Py_XDECREF(izip);
Py_DECREF(ret); return result;
return NULL;
} }
PyDoc_STRVAR(zip_doc, PyDoc_STRVAR(zip_doc,
"zip(seq1 [, seq2 [...]]) -> [(seq1[0], seq2[0] ...), (...)]\n\ "zip(it1 [, it2 [...]]) -> iter([(it1[0], it2[0] ...), ...])\n\
\n\ \n\
Return a list of tuples, where each tuple contains the i-th element\n\ Return an iterator yielding tuples, where each tuple contains the\n\
from each of the argument sequences. The returned list is truncated\n\ corresponding element from each of the argument iterables.\n\
in length to the length of the shortest argument sequence."); The returned iterator ends when the shortest argument iterable is exhausted.\n\
NOTE: This is implemented using itertools.izip().");
static PyMethodDef builtin_methods[] = { static PyMethodDef builtin_methods[] = {
......
Markdown is supported
0%
or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment