[Python-checkins] [2.7] bpo-30347: Stop crashes when concurrently iterate over itertools.groupby() iterators. (GH-1557). (#3772)

[Serhiy Storchaka]

https://github.com/python/cpython/commit/d0b9dc33676bdad217d5074954c1b37d4ca54a10
commit: d0b9dc33676bdad217d5074954c1b37d4ca54a10
branch: 2.7
author: Serhiy Storchaka <storchaka at gmail.com>
committer: GitHub <noreply at github.com>
date: 2017-09-26T23:15:36+03:00
summary:

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

(cherry picked from commit c740e4fe8a9bc5815dc18c38d7f7600b128c3c51)

files:
A Misc/NEWS.d/next/Library/2017-09-25-14-04-30.bpo-30347.B4--_D.rst
M Lib/test/test_itertools.py
M Modules/itertoolsmodule.c

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 @@ 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(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);