Fix exception handling in itertools.izip_longest().

This commit is contained in:
Raymond Hettinger 2009-11-01 21:02:38 +00:00
parent 6c4a725a37
commit a9311a3c50
3 changed files with 79 additions and 34 deletions

View File

@ -581,6 +581,46 @@ class TestBasicOps(unittest.TestCase):
ids = list(map(id, list(zip_longest('abc', 'def')))) ids = list(map(id, list(zip_longest('abc', 'def'))))
self.assertEqual(len(dict.fromkeys(ids)), len(ids)) self.assertEqual(len(dict.fromkeys(ids)), len(ids))
def test_bug_7244(self):
class Repeater:
# this class is similar to itertools.repeat
def __init__(self, o, t, e):
self.o = o
self.t = int(t)
self.e = e
def __iter__(self): # its iterator is itself
return self
def __next__(self):
if self.t > 0:
self.t -= 1
return self.o
else:
raise self.e
# Formerly this code in would fail in debug mode
# with Undetected Error and Stop Iteration
r1 = Repeater(1, 3, StopIteration)
r2 = Repeater(2, 4, StopIteration)
def run(r1, r2):
result = []
for i, j in zip_longest(r1, r2, fillvalue=0):
with support.captured_output('stdout'):
print((i, j))
result.append((i, j))
return result
self.assertEqual(run(r1, r2), [(1,2), (1,2), (1,2), (0,2)])
# Formerly, the RuntimeError would be lost
# and StopIteration would stop as expected
r1 = Repeater(1, 3, RuntimeError)
r2 = Repeater(2, 4, StopIteration)
it = zip_longest(r1, r2, fillvalue=0)
self.assertEqual(next(it), (1, 2))
self.assertEqual(next(it), (1, 2))
self.assertEqual(next(it), (1, 2))
self.assertRaises(RuntimeError, next, it)
def test_product(self): def test_product(self):
for args, result in [ for args, result in [
([], [()]), # zero iterables ([], [()]), # zero iterables

View File

@ -12,6 +12,9 @@ What's New in Python 3.1.2?
Core and Builtins Core and Builtins
----------------- -----------------
- Issue #7244: itertools.izip_longest() no longer ignores exceptions
raised during the formation of an output tuple.
- Issue #3297: On wide unicode builds, do not split unicode characters into - Issue #3297: On wide unicode builds, do not split unicode characters into
surrogates. surrogates.

View File

@ -3344,71 +3344,73 @@ zip_longest_traverse(ziplongestobject *lz, visitproc visit, void *arg)
static PyObject * static PyObject *
zip_longest_next(ziplongestobject *lz) zip_longest_next(ziplongestobject *lz)
{ {
Py_ssize_t i; Py_ssize_t i;
Py_ssize_t tuplesize = lz->tuplesize; Py_ssize_t tuplesize = lz->tuplesize;
PyObject *result = lz->result; PyObject *result = lz->result;
PyObject *it; PyObject *it;
PyObject *item; PyObject *item;
PyObject *olditem; PyObject *olditem;
if (tuplesize == 0) if (tuplesize == 0)
return NULL; return NULL;
if (lz->numactive == 0) if (lz->numactive == 0)
return NULL; return NULL;
if (Py_REFCNT(result) == 1) { if (Py_REFCNT(result) == 1) {
Py_INCREF(result); Py_INCREF(result);
for (i=0 ; i < tuplesize ; i++) { for (i=0 ; i < tuplesize ; i++) {
it = PyTuple_GET_ITEM(lz->ittuple, i); it = PyTuple_GET_ITEM(lz->ittuple, i);
if (it == NULL) { if (it == NULL) {
Py_INCREF(lz->fillvalue); Py_INCREF(lz->fillvalue);
item = lz->fillvalue; item = lz->fillvalue;
} else { } else {
item = (*Py_TYPE(it)->tp_iternext)(it); item = PyIter_Next(it);
if (item == NULL) { if (item == NULL) {
lz->numactive -= 1; lz->numactive -= 1;
if (lz->numactive == 0) { if (lz->numactive == 0 || PyErr_Occurred()) {
lz->numactive = 0;
Py_DECREF(result); Py_DECREF(result);
return NULL; return NULL;
} else { } else {
Py_INCREF(lz->fillvalue); Py_INCREF(lz->fillvalue);
item = lz->fillvalue; item = lz->fillvalue;
PyTuple_SET_ITEM(lz->ittuple, i, NULL); PyTuple_SET_ITEM(lz->ittuple, i, NULL);
Py_DECREF(it); Py_DECREF(it);
} }
} }
} }
olditem = PyTuple_GET_ITEM(result, i); olditem = PyTuple_GET_ITEM(result, i);
PyTuple_SET_ITEM(result, i, item); PyTuple_SET_ITEM(result, i, item);
Py_DECREF(olditem); Py_DECREF(olditem);
} }
} else { } else {
result = PyTuple_New(tuplesize); result = PyTuple_New(tuplesize);
if (result == NULL) if (result == NULL)
return NULL; return NULL;
for (i=0 ; i < tuplesize ; i++) { for (i=0 ; i < tuplesize ; i++) {
it = PyTuple_GET_ITEM(lz->ittuple, i); it = PyTuple_GET_ITEM(lz->ittuple, i);
if (it == NULL) { if (it == NULL) {
Py_INCREF(lz->fillvalue); Py_INCREF(lz->fillvalue);
item = lz->fillvalue; item = lz->fillvalue;
} else { } else {
item = (*Py_TYPE(it)->tp_iternext)(it); item = PyIter_Next(it);
if (item == NULL) { if (item == NULL) {
lz->numactive -= 1; lz->numactive -= 1;
if (lz->numactive == 0) { if (lz->numactive == 0 || PyErr_Occurred()) {
lz->numactive = 0;
Py_DECREF(result); Py_DECREF(result);
return NULL; return NULL;
} else { } else {
Py_INCREF(lz->fillvalue); Py_INCREF(lz->fillvalue);
item = lz->fillvalue; item = lz->fillvalue;
PyTuple_SET_ITEM(lz->ittuple, i, NULL); PyTuple_SET_ITEM(lz->ittuple, i, NULL);
Py_DECREF(it); Py_DECREF(it);
} }
} }
} }
PyTuple_SET_ITEM(result, i, item); PyTuple_SET_ITEM(result, i, item);
} }
} }
return result; return result;
} }
PyDoc_STRVAR(zip_longest_doc, PyDoc_STRVAR(zip_longest_doc,