bpo-30346: An iterator produced by the itertools.groupby() iterator (#1569)

now becames exhausted after advancing the groupby iterator.

4 files changed, 34 insertions, 3 deletions

diff --git a/Doc/library/itertools.rst b/Doc/library/itertools.rst
index c989e46..530c29d 100644
--- a/Doc/library/itertools.rst
+++ b/Doc/library/itertools.rst
@@ -401,13 +401,14 @@ loops that truncate the stream.
           def __iter__(self):
               return self
           def __next__(self):
+              self.id = object()
               while self.currkey == self.tgtkey:
                   self.currvalue = next(self.it)    # Exit on StopIteration
                   self.currkey = self.keyfunc(self.currvalue)
               self.tgtkey = self.currkey
-              return (self.currkey, self._grouper(self.tgtkey))
-          def _grouper(self, tgtkey):
-              while self.currkey == tgtkey:
+              return (self.currkey, self._grouper(self.tgtkey, self.id))
+          def _grouper(self, tgtkey, id):
+              while self.id is id and self.currkey == tgtkey:
                   yield self.currvalue
                   try:
                       self.currvalue = next(self.it)
diff --git a/Lib/test/test_itertools.py b/Lib/test/test_itertools.py
index 50cf148..8353e68 100644
--- a/Lib/test/test_itertools.py
+++ b/Lib/test/test_itertools.py
@@ -751,6 +751,26 @@ class TestBasicOps(unittest.TestCase):
         self.assertEqual(set(keys), expectedkeys)
         self.assertEqual(len(keys), len(expectedkeys))
 
+        # Check case where inner iterator is used after advancing the groupby
+        # iterator
+        s = list(zip('AABBBAAAA', range(9)))
+        it = groupby(s, testR)
+        _, g1 = next(it)
+        _, g2 = next(it)
+        _, g3 = next(it)
+        self.assertEqual(list(g1), [])
+        self.assertEqual(list(g2), [])
+        self.assertEqual(next(g3), ('A', 5))
+        list(it)  # exhaust the groupby iterator
+        self.assertEqual(list(g3), [])
+
+        for proto in range(pickle.HIGHEST_PROTOCOL + 1):
+            it = groupby(s, testR)
+            _, g = next(it)
+            next(it)
+            next(it)
+            self.assertEqual(list(pickle.loads(pickle.dumps(g, proto))), [])
+
         # Exercise pipes and filters style
         s = 'abracadabra'
         # sort s | uniq
diff --git a/Misc/NEWS.d/next/Library/2017-09-24-13-08-46.bpo-30346.Csse77.rst b/Misc/NEWS.d/next/Library/2017-09-24-13-08-46.bpo-30346.Csse77.rst
new file mode 100644
index 0000000..81ad053
--- /dev/null
+++ b/Misc/NEWS.d/next/Library/2017-09-24-13-08-46.bpo-30346.Csse77.rst
@@ -0,0 +1,2 @@
+An iterator produced by itertools.groupby() iterator now becames exhausted
+after advancing the groupby iterator.
diff --git a/Modules/itertoolsmodule.c b/Modules/itertoolsmodule.c
index 48e6c35..2ac5ab2 100644
--- a/Modules/itertoolsmodule.c
+++ b/Modules/itertoolsmodule.c
@@ -17,6 +17,7 @@ typedef struct {
     PyObject *tgtkey;
     PyObject *currkey;
     PyObject *currvalue;
+    const void *currgrouper;  /* borrowed reference */
 } groupbyobject;
 
 static PyTypeObject groupby_type;
@@ -77,6 +78,7 @@ groupby_next(groupbyobject *gbo)
 {
     PyObject *newvalue, *newkey, *r, *grouper;
 
+    gbo->currgrouper = NULL;
     /* skip to next iteration group */
     for (;;) {
         if (gbo->currkey == NULL)
@@ -255,6 +257,7 @@ _grouper_create(groupbyobject *parent, PyObject *tgtkey)
     Py_INCREF(parent);
     igo->tgtkey = tgtkey;
     Py_INCREF(tgtkey);
+    parent->currgrouper = igo;  /* borrowed reference */
 
     PyObject_GC_Track(igo);
     return (PyObject *)igo;
@@ -284,6 +287,8 @@ _grouper_next(_grouperobject *igo)
     PyObject *newvalue, *newkey, *r;
     int rcmp;
 
+    if (gbo->currgrouper != igo)
+        return NULL;
     if (gbo->currvalue == NULL) {
         newvalue = PyIter_Next(gbo->it);
         if (newvalue == NULL)
@@ -321,6 +326,9 @@ _grouper_next(_grouperobject *igo)
 static PyObject *
 _grouper_reduce(_grouperobject *lz)
 {
+    if (((groupbyobject *)lz->parent)->currgrouper != lz) {
+        return Py_BuildValue("N(())", _PyObject_GetBuiltin("iter"));
+    }
     return Py_BuildValue("O(OO)", Py_TYPE(lz), lz->parent, lz->tgtkey);
 }