bpo-30346: An iterator produced by the itertools.groupby() iterator (… · python/cpython@c247caf (original) (raw)

4 files changed

lines changed

Original file line number Diff line number Diff line change
@@ -401,13 +401,14 @@ loops that truncate the stream.
401 401 def __iter__(self):
402 402 return self
403 403 def __next__(self):
404 + self.id = object()
404 405 while self.currkey == self.tgtkey:
405 406 self.currvalue = next(self.it) # Exit on StopIteration
406 407 self.currkey = self.keyfunc(self.currvalue)
407 408 self.tgtkey = self.currkey
408 - return (self.currkey, self._grouper(self.tgtkey))
409 - def _grouper(self, tgtkey):
410 - while self.currkey == tgtkey:
409 + return (self.currkey, self._grouper(self.tgtkey, self.id))
410 + def _grouper(self, tgtkey, id):
411 + while self.id is id and self.currkey == tgtkey:
411 412 yield self.currvalue
412 413 try:
413 414 self.currvalue = next(self.it)
Original file line number Diff line number Diff line change
@@ -751,6 +751,26 @@ def test_groupby(self):
751 751 self.assertEqual(set(keys), expectedkeys)
752 752 self.assertEqual(len(keys), len(expectedkeys))
753 753
754 +# Check case where inner iterator is used after advancing the groupby
755 +# iterator
756 +s = list(zip('AABBBAAAA', range(9)))
757 +it = groupby(s, testR)
758 +_, g1 = next(it)
759 +_, g2 = next(it)
760 +_, g3 = next(it)
761 +self.assertEqual(list(g1), [])
762 +self.assertEqual(list(g2), [])
763 +self.assertEqual(next(g3), ('A', 5))
764 +list(it) # exhaust the groupby iterator
765 +self.assertEqual(list(g3), [])
766 +
767 +for proto in range(pickle.HIGHEST_PROTOCOL + 1):
768 +it = groupby(s, testR)
769 +_, g = next(it)
770 +next(it)
771 +next(it)
772 +self.assertEqual(list(pickle.loads(pickle.dumps(g, proto))), [])
773 +
754 774 # Exercise pipes and filters style
755 775 s = 'abracadabra'
756 776 # sort s | uniq
Original file line number Diff line number Diff line change
@@ -0,0 +1,2 @@
1 +An iterator produced by itertools.groupby() iterator now becames exhausted
2 +after advancing the groupby iterator.
Original file line number Diff line number Diff line change
@@ -17,6 +17,7 @@ typedef struct {
17 17 PyObject *tgtkey;
18 18 PyObject *currkey;
19 19 PyObject *currvalue;
20 +const void *currgrouper; /* borrowed reference */
20 21 } groupbyobject;
21 22
22 23 static PyTypeObject groupby_type;
@@ -77,6 +78,7 @@ groupby_next(groupbyobject *gbo)
77 78 {
78 79 PyObject *newvalue, *newkey, *r, *grouper;
79 80
81 +gbo->currgrouper = NULL;
80 82 /* skip to next iteration group */
81 83 for (;;) {
82 84 if (gbo->currkey == NULL)
@@ -255,6 +257,7 @@ _grouper_create(groupbyobject *parent, PyObject *tgtkey)
255 257 Py_INCREF(parent);
256 258 igo->tgtkey = tgtkey;
257 259 Py_INCREF(tgtkey);
260 +parent->currgrouper = igo; /* borrowed reference */
258 261
259 262 PyObject_GC_Track(igo);
260 263 return (PyObject *)igo;
@@ -284,6 +287,8 @@ _grouper_next(_grouperobject *igo)
284 287 PyObject *newvalue, *newkey, *r;
285 288 int rcmp;
286 289
290 +if (gbo->currgrouper != igo)
291 +return NULL;
287 292 if (gbo->currvalue == NULL) {
288 293 newvalue = PyIter_Next(gbo->it);
289 294 if (newvalue == NULL)
@@ -321,6 +326,9 @@ _grouper_next(_grouperobject *igo)
321 326 static PyObject *
322 327 _grouper_reduce(_grouperobject *lz)
323 328 {
329 +if (((groupbyobject *)lz->parent)->currgrouper != lz) {
330 +return Py_BuildValue("N(())", _PyObject_GetBuiltin("iter"));
331 + }
324 332 return Py_BuildValue("O(OO)", Py_TYPE(lz), lz->parent, lz->tgtkey);
325 333 }
326 334