[3.6] bpo-30347: Stop crashes when concurrently iterate over itertools.groupby() iterators. (GH-1557) (#3770)

(cherry picked from commit c740e4f)

Author: miss-islington

Lib/test/test_itertools.py

@@ -1984,6 +1984,30 @@ def test_long_chain_of_empty_iterables(self):
  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...

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.

Modules/itertoolsmodule.c

@@ -72,10 +72,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;
 
  /* skip to next iteration group */
  for (;;) {
@@ -93,25 +120,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);
 
@@ -282,28 +293,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);