Skip to content

Commit 69b2dc8

Browse files
miss-islingtonserhiy-storchaka
authored andcommitted
[3.6] bpo-30347: Stop crashes when concurrently iterate over itertools.groupby() iterators. (GH-1557) (#3770)
(cherry picked from commit c740e4f)
1 parent d6a3562 commit 69b2dc8

File tree

3 files changed

+56
-36
lines changed

3 files changed

+56
-36
lines changed

Lib/test/test_itertools.py

Lines changed: 24 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1984,6 +1984,30 @@ def test_long_chain_of_empty_iterables(self):
19841984
with self.assertRaises(StopIteration):
19851985
next(it)
19861986

1987+
def test_issue30347_1(self):
1988+
def f(n):
1989+
if n == 5:
1990+
list(b)
1991+
return n != 6
1992+
for (k, b) in groupby(range(10), f):
1993+
list(b) # shouldn't crash
1994+
1995+
def test_issue30347_2(self):
1996+
class K:
1997+
def __init__(self, v):
1998+
pass
1999+
def __eq__(self, other):
2000+
nonlocal i
2001+
i += 1
2002+
if i == 1:
2003+
next(g, None)
2004+
return True
2005+
i = 0
2006+
g = next(groupby(range(10), K))[1]
2007+
for j in range(2):
2008+
next(g, None) # shouldn't crash
2009+
2010+
19872011
class SubclassWithKwargsTest(unittest.TestCase):
19882012
def test_keywords_in_subclass(self):
19892013
# count is not subclassable...
Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1 @@
1+
Stop crashes when concurrently iterate over itertools.groupby() iterators.

Modules/itertoolsmodule.c

Lines changed: 31 additions & 36 deletions
Original file line numberDiff line numberDiff line change
@@ -72,10 +72,37 @@ groupby_traverse(groupbyobject *gbo, visitproc visit, void *arg)
7272
return 0;
7373
}
7474

75+
Py_LOCAL_INLINE(int)
76+
groupby_step(groupbyobject *gbo)
77+
{
78+
PyObject *newvalue, *newkey, *oldvalue;
79+
80+
newvalue = PyIter_Next(gbo->it);
81+
if (newvalue == NULL)
82+
return -1;
83+
84+
if (gbo->keyfunc == Py_None) {
85+
newkey = newvalue;
86+
Py_INCREF(newvalue);
87+
} else {
88+
newkey = PyObject_CallFunctionObjArgs(gbo->keyfunc, newvalue, NULL);
89+
if (newkey == NULL) {
90+
Py_DECREF(newvalue);
91+
return -1;
92+
}
93+
}
94+
95+
oldvalue = gbo->currvalue;
96+
gbo->currvalue = newvalue;
97+
Py_XSETREF(gbo->currkey, newkey);
98+
Py_XDECREF(oldvalue);
99+
return 0;
100+
}
101+
75102
static PyObject *
76103
groupby_next(groupbyobject *gbo)
77104
{
78-
PyObject *newvalue, *newkey, *r, *grouper;
105+
PyObject *r, *grouper;
79106

80107
/* skip to next iteration group */
81108
for (;;) {
@@ -93,25 +120,9 @@ groupby_next(groupbyobject *gbo)
93120
break;
94121
}
95122

96-
newvalue = PyIter_Next(gbo->it);
97-
if (newvalue == NULL)
123+
if (groupby_step(gbo) < 0)
98124
return NULL;
99-
100-
if (gbo->keyfunc == Py_None) {
101-
newkey = newvalue;
102-
Py_INCREF(newvalue);
103-
} else {
104-
newkey = PyObject_CallFunctionObjArgs(gbo->keyfunc, newvalue, NULL);
105-
if (newkey == NULL) {
106-
Py_DECREF(newvalue);
107-
return NULL;
108-
}
109-
}
110-
111-
Py_XSETREF(gbo->currkey, newkey);
112-
Py_XSETREF(gbo->currvalue, newvalue);
113125
}
114-
115126
Py_INCREF(gbo->currkey);
116127
Py_XSETREF(gbo->tgtkey, gbo->currkey);
117128

@@ -282,28 +293,12 @@ static PyObject *
282293
_grouper_next(_grouperobject *igo)
283294
{
284295
groupbyobject *gbo = (groupbyobject *)igo->parent;
285-
PyObject *newvalue, *newkey, *r;
296+
PyObject *r;
286297
int rcmp;
287298

288299
if (gbo->currvalue == NULL) {
289-
newvalue = PyIter_Next(gbo->it);
290-
if (newvalue == NULL)
300+
if (groupby_step(gbo) < 0)
291301
return NULL;
292-
293-
if (gbo->keyfunc == Py_None) {
294-
newkey = newvalue;
295-
Py_INCREF(newvalue);
296-
} else {
297-
newkey = PyObject_CallFunctionObjArgs(gbo->keyfunc, newvalue, NULL);
298-
if (newkey == NULL) {
299-
Py_DECREF(newvalue);
300-
return NULL;
301-
}
302-
}
303-
304-
assert(gbo->currkey == NULL);
305-
gbo->currkey = newkey;
306-
gbo->currvalue = newvalue;
307302
}
308303

309304
assert(gbo->currkey != NULL);

0 commit comments

Comments
 (0)