diff --git a/Lib/test/test_b2.py b/Lib/test/test_b2.py index 28022153b56..324d02fdd1a 100644 --- 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 a50c74f2fd5..ddc58a7cfc7 100644 --- 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 aecc5e9114c..1f971cd40c9 100644 --- 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 4a51ccd98a1..cc1bc957a96 100644 --- 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; }