diff --git a/Lib/test/test_itertools.py b/Lib/test/test_itertools.py index 7b2362d29e9..7eee81a568c 100644 --- a/Lib/test/test_itertools.py +++ b/Lib/test/test_itertools.py @@ -1473,6 +1473,29 @@ class RegressionTests(unittest.TestCase): with self.assertRaises(StopIteration): next(it) + def test_issue30347_1(self): + def f(n): + if n == 5: + list(b) + return n != 6 + for (k, b) in groupby(range(10), f): + list(b) # shouldn't crash + + def test_issue30347_2(self): + class K(object): + i = 0 + def __init__(self, v): + pass + def __eq__(self, other): + K.i += 1 + if K.i == 1: + next(g, None) + return True + g = next(groupby(range(10), K))[1] + for j in range(2): + next(g, None) # shouldn't crash + + class SubclassWithKwargsTest(unittest.TestCase): def test_keywords_in_subclass(self): # count is not subclassable... diff --git a/Misc/NEWS.d/next/Library/2017-09-25-14-04-30.bpo-30347.B4--_D.rst b/Misc/NEWS.d/next/Library/2017-09-25-14-04-30.bpo-30347.B4--_D.rst new file mode 100644 index 00000000000..859c6415529 --- /dev/null +++ b/Misc/NEWS.d/next/Library/2017-09-25-14-04-30.bpo-30347.B4--_D.rst @@ -0,0 +1 @@ +Stop crashes when concurrently iterate over itertools.groupby() iterators. diff --git a/Modules/itertoolsmodule.c b/Modules/itertoolsmodule.c index cf3aadfd3e0..47db7affb05 100644 --- a/Modules/itertoolsmodule.c +++ b/Modules/itertoolsmodule.c @@ -71,10 +71,37 @@ groupby_traverse(groupbyobject *gbo, visitproc visit, void *arg) return 0; } +Py_LOCAL_INLINE(int) +groupby_step(groupbyobject *gbo) +{ + PyObject *newvalue, *newkey, *oldvalue; + + newvalue = PyIter_Next(gbo->it); + if (newvalue == NULL) + return -1; + + if (gbo->keyfunc == Py_None) { + newkey = newvalue; + Py_INCREF(newvalue); + } else { + newkey = PyObject_CallFunctionObjArgs(gbo->keyfunc, newvalue, NULL); + if (newkey == NULL) { + Py_DECREF(newvalue); + return -1; + } + } + + oldvalue = gbo->currvalue; + gbo->currvalue = newvalue; + Py_XSETREF(gbo->currkey, newkey); + Py_XDECREF(oldvalue); + return 0; +} + static PyObject * groupby_next(groupbyobject *gbo) { - PyObject *newvalue, *newkey, *r, *grouper, *tmp; + PyObject *r, *grouper; /* skip to next iteration group */ for (;;) { @@ -93,35 +120,11 @@ groupby_next(groupbyobject *gbo) break; } - newvalue = PyIter_Next(gbo->it); - if (newvalue == NULL) + if (groupby_step(gbo) < 0) return NULL; - - if (gbo->keyfunc == Py_None) { - newkey = newvalue; - Py_INCREF(newvalue); - } else { - newkey = PyObject_CallFunctionObjArgs(gbo->keyfunc, - newvalue, NULL); - if (newkey == NULL) { - Py_DECREF(newvalue); - return NULL; - } - } - - tmp = gbo->currkey; - gbo->currkey = newkey; - Py_XDECREF(tmp); - - tmp = gbo->currvalue; - gbo->currvalue = newvalue; - Py_XDECREF(tmp); } - Py_INCREF(gbo->currkey); - tmp = gbo->tgtkey; - gbo->tgtkey = gbo->currkey; - Py_XDECREF(tmp); + Py_XSETREF(gbo->tgtkey, gbo->currkey); grouper = _grouper_create(gbo, gbo->tgtkey); if (grouper == NULL) @@ -229,29 +232,12 @@ static PyObject * _grouper_next(_grouperobject *igo) { groupbyobject *gbo = (groupbyobject *)igo->parent; - PyObject *newvalue, *newkey, *r; + PyObject *r; int rcmp; if (gbo->currvalue == NULL) { - newvalue = PyIter_Next(gbo->it); - if (newvalue == NULL) + if (groupby_step(gbo) < 0) return NULL; - - if (gbo->keyfunc == Py_None) { - newkey = newvalue; - Py_INCREF(newvalue); - } else { - newkey = PyObject_CallFunctionObjArgs(gbo->keyfunc, - newvalue, NULL); - if (newkey == NULL) { - Py_DECREF(newvalue); - return NULL; - } - } - - assert(gbo->currkey == NULL); - gbo->currkey = newkey; - gbo->currvalue = newvalue; } assert(gbo->currkey != NULL);