Skip to content

Commit e3d380d

Browse files
[3.11] gh-109786: Fix leaks and crash when re-enter itertools.pairwise.__next__() (GH-109788) (GH-112700)
(cherry picked from commit 6ca9d3e) Co-authored-by: Serhiy Storchaka <[email protected]>
1 parent 28afd8d commit e3d380d

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
@@ -1079,6 +1079,78 @@ def test_pairwise(self):
10791079
with self.assertRaises(TypeError):
10801080
pairwise(None) # non-iterable argument
10811081

1082+
def test_pairwise_reenter(self):
1083+
def check(reenter_at, expected):
1084+
class I:
1085+
count = 0
1086+
def __iter__(self):
1087+
return self
1088+
def __next__(self):
1089+
self.count +=1
1090+
if self.count in reenter_at:
1091+
return next(it)
1092+
return [self.count] # new object
1093+
1094+
it = pairwise(I())
1095+
for item in expected:
1096+
self.assertEqual(next(it), item)
1097+
1098+
check({1}, [
1099+
(([2], [3]), [4]),
1100+
([4], [5]),
1101+
])
1102+
check({2}, [
1103+
([1], ([1], [3])),
1104+
(([1], [3]), [4]),
1105+
([4], [5]),
1106+
])
1107+
check({3}, [
1108+
([1], [2]),
1109+
([2], ([2], [4])),
1110+
(([2], [4]), [5]),
1111+
([5], [6]),
1112+
])
1113+
check({1, 2}, [
1114+
((([3], [4]), [5]), [6]),
1115+
([6], [7]),
1116+
])
1117+
check({1, 3}, [
1118+
(([2], ([2], [4])), [5]),
1119+
([5], [6]),
1120+
])
1121+
check({1, 4}, [
1122+
(([2], [3]), (([2], [3]), [5])),
1123+
((([2], [3]), [5]), [6]),
1124+
([6], [7]),
1125+
])
1126+
check({2, 3}, [
1127+
([1], ([1], ([1], [4]))),
1128+
(([1], ([1], [4])), [5]),
1129+
([5], [6]),
1130+
])
1131+
1132+
def test_pairwise_reenter2(self):
1133+
def check(maxcount, expected):
1134+
class I:
1135+
count = 0
1136+
def __iter__(self):
1137+
return self
1138+
def __next__(self):
1139+
if self.count >= maxcount:
1140+
raise StopIteration
1141+
self.count +=1
1142+
if self.count == 1:
1143+
return next(it, None)
1144+
return [self.count] # new object
1145+
1146+
it = pairwise(I())
1147+
self.assertEqual(list(it), expected)
1148+
1149+
check(1, [])
1150+
check(2, [])
1151+
check(3, [])
1152+
check(4, [(([2], [3]), [4])])
1153+
10821154
def test_product(self):
10831155
for args, result in [
10841156
([], [()]), # 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
@@ -119,21 +119,30 @@ pairwise_next(pairwiseobject *po)
119119
return NULL;
120120
}
121121
if (old == NULL) {
122-
po->old = old = (*Py_TYPE(it)->tp_iternext)(it);
122+
old = (*Py_TYPE(it)->tp_iternext)(it);
123+
Py_XSETREF(po->old, old);
123124
if (old == NULL) {
124125
Py_CLEAR(po->it);
125126
return NULL;
126127
}
128+
it = po->it;
129+
if (it == NULL) {
130+
Py_CLEAR(po->old);
131+
return NULL;
132+
}
127133
}
134+
Py_INCREF(old);
128135
new = (*Py_TYPE(it)->tp_iternext)(it);
129136
if (new == NULL) {
130137
Py_CLEAR(po->it);
131138
Py_CLEAR(po->old);
139+
Py_DECREF(old);
132140
return NULL;
133141
}
134142
/* Future optimization: Reuse the result tuple as we do in enumerate() */
135143
result = PyTuple_Pack(2, old, new);
136-
Py_SETREF(po->old, new);
144+
Py_XSETREF(po->old, new);
145+
Py_DECREF(old);
137146
return result;
138147
}
139148

0 commit comments

Comments
 (0)