diff options
| -rw-r--r-- | Lib/test/test_itertools.py | 24 | ||||
| -rw-r--r-- | Misc/NEWS.d/next/Library/2017-09-25-14-04-30.bpo-30347.B4--_D.rst | 1 | ||||
| -rw-r--r-- | Modules/itertoolsmodule.c | 67 |
3 files changed, 56 insertions, 36 deletions
diff --git a/Lib/test/test_itertools.py b/Lib/test/test_itertools.py index 8353e68..a359905 100644 --- a/Lib/test/test_itertools.py +++ b/Lib/test/test_itertools.py @@ -2017,6 +2017,30 @@ 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: + def __init__(self, v): + pass + def __eq__(self, other): + nonlocal i + i += 1 + if i == 1: + next(g, None) + return True + i = 0 + 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 0000000..859c641 --- /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 cc9895a..9190236 100644 --- a/Modules/itertoolsmodule.c +++ b/Modules/itertoolsmodule.c @@ -73,10 +73,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; + PyObject *r, *grouper; gbo->currgrouper = NULL; /* skip to next iteration group */ @@ -95,25 +122,9 @@ 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; - } - } - - Py_XSETREF(gbo->currkey, newkey); - Py_XSETREF(gbo->currvalue, newvalue); } - Py_INCREF(gbo->currkey); Py_XSETREF(gbo->tgtkey, gbo->currkey); @@ -285,30 +296,14 @@ static PyObject * _grouper_next(_grouperobject *igo) { groupbyobject *gbo = (groupbyobject *)igo->parent; - PyObject *newvalue, *newkey, *r; + PyObject *r; int rcmp; if (gbo->currgrouper != igo) return NULL; 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); |