Skip to content

Commit 6ca9d3e

Browse files
gh-109786: Fix leaks and crash when re-enter itertools.pairwise.__next__() (GH-109788)
1 parent c74e9fb commit 6ca9d3e

File tree

3 files changed

+85
-2
lines changed

3 files changed

+85
-2
lines changed

Lib/test/test_itertools.py

Lines changed: 72 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1152,6 +1152,78 @@ def test_pairwise(self):
11521152
with self.assertRaises(TypeError):
11531153
pairwise(None) # non-iterable argument
11541154

1155+
def test_pairwise_reenter(self):
1156+
def check(reenter_at, expected):
1157+
class I:
1158+
count = 0
1159+
def __iter__(self):
1160+
return self
1161+
def __next__(self):
1162+
self.count +=1
1163+
if self.count in reenter_at:
1164+
return next(it)
1165+
return [self.count] # new object
1166+
1167+
it = pairwise(I())
1168+
for item in expected:
1169+
self.assertEqual(next(it), item)
1170+
1171+
check({1}, [
1172+
(([2], [3]), [4]),
1173+
([4], [5]),
1174+
])
1175+
check({2}, [
1176+
([1], ([1], [3])),
1177+
(([1], [3]), [4]),
1178+
([4], [5]),
1179+
])
1180+
check({3}, [
1181+
([1], [2]),
1182+
([2], ([2], [4])),
1183+
(([2], [4]), [5]),
1184+
([5], [6]),
1185+
])
1186+
check({1, 2}, [
1187+
((([3], [4]), [5]), [6]),
1188+
([6], [7]),
1189+
])
1190+
check({1, 3}, [
1191+
(([2], ([2], [4])), [5]),
1192+
([5], [6]),
1193+
])
1194+
check({1, 4}, [
1195+
(([2], [3]), (([2], [3]), [5])),
1196+
((([2], [3]), [5]), [6]),
1197+
([6], [7]),
1198+
])
1199+
check({2, 3}, [
1200+
([1], ([1], ([1], [4]))),
1201+
(([1], ([1], [4])), [5]),
1202+
([5], [6]),
1203+
])
1204+
1205+
def test_pairwise_reenter2(self):
1206+
def check(maxcount, expected):
1207+
class I:
1208+
count = 0
1209+
def __iter__(self):
1210+
return self
1211+
def __next__(self):
1212+
if self.count >= maxcount:
1213+
raise StopIteration
1214+
self.count +=1
1215+
if self.count == 1:
1216+
return next(it, None)
1217+
return [self.count] # new object
1218+
1219+
it = pairwise(I())
1220+
self.assertEqual(list(it), expected)
1221+
1222+
check(1, [])
1223+
check(2, [])
1224+
check(3, [])
1225+
check(4, [(([2], [3]), [4])])
1226+
11551227
def test_product(self):
11561228
for args, result in [
11571229
([], [()]), # zero iterables
Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,2 @@
1+
Fix possible reference leaks and crash when re-enter the ``__next__()`` method of
2+
:class:`itertools.pairwise`.

Modules/itertoolsmodule.c

Lines changed: 11 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -330,21 +330,30 @@ pairwise_next(pairwiseobject *po)
330330
return NULL;
331331
}
332332
if (old == NULL) {
333-
po->old = old = (*Py_TYPE(it)->tp_iternext)(it);
333+
old = (*Py_TYPE(it)->tp_iternext)(it);
334+
Py_XSETREF(po->old, old);
334335
if (old == NULL) {
335336
Py_CLEAR(po->it);
336337
return NULL;
337338
}
339+
it = po->it;
340+
if (it == NULL) {
341+
Py_CLEAR(po->old);
342+
return NULL;
343+
}
338344
}
345+
Py_INCREF(old);
339346
new = (*Py_TYPE(it)->tp_iternext)(it);
340347
if (new == NULL) {
341348
Py_CLEAR(po->it);
342349
Py_CLEAR(po->old);
350+
Py_DECREF(old);
343351
return NULL;
344352
}
345353
/* Future optimization: Reuse the result tuple as we do in enumerate() */
346354
result = PyTuple_Pack(2, old, new);
347-
Py_SETREF(po->old, new);
355+
Py_XSETREF(po->old, new);
356+
Py_DECREF(old);
348357
return result;
349358
}
350359

0 commit comments

Comments
 (0)