Skip to content

Commit

Permalink
pythongh-109786: Fix leaks and crash when re-enter itertools.pairwise…
Browse files Browse the repository at this point in the history
….__next__() (pythonGH-109788)
  • Loading branch information
serhiy-storchaka authored and Glyphack committed Jan 27, 2024
1 parent 78a88b7 commit ea831c5
Show file tree
Hide file tree
Showing 3 changed files with 85 additions and 2 deletions.
72 changes: 72 additions & 0 deletions Lib/test/test_itertools.py
Original file line number Diff line number Diff line change
Expand Up @@ -1152,6 +1152,78 @@ def test_pairwise(self):
with self.assertRaises(TypeError):
pairwise(None) # non-iterable argument

def test_pairwise_reenter(self):
def check(reenter_at, expected):
class I:
count = 0
def __iter__(self):
return self
def __next__(self):
self.count +=1
if self.count in reenter_at:
return next(it)
return [self.count] # new object

it = pairwise(I())
for item in expected:
self.assertEqual(next(it), item)

check({1}, [
(([2], [3]), [4]),
([4], [5]),
])
check({2}, [
([1], ([1], [3])),
(([1], [3]), [4]),
([4], [5]),
])
check({3}, [
([1], [2]),
([2], ([2], [4])),
(([2], [4]), [5]),
([5], [6]),
])
check({1, 2}, [
((([3], [4]), [5]), [6]),
([6], [7]),
])
check({1, 3}, [
(([2], ([2], [4])), [5]),
([5], [6]),
])
check({1, 4}, [
(([2], [3]), (([2], [3]), [5])),
((([2], [3]), [5]), [6]),
([6], [7]),
])
check({2, 3}, [
([1], ([1], ([1], [4]))),
(([1], ([1], [4])), [5]),
([5], [6]),
])

def test_pairwise_reenter2(self):
def check(maxcount, expected):
class I:
count = 0
def __iter__(self):
return self
def __next__(self):
if self.count >= maxcount:
raise StopIteration
self.count +=1
if self.count == 1:
return next(it, None)
return [self.count] # new object

it = pairwise(I())
self.assertEqual(list(it), expected)

check(1, [])
check(2, [])
check(3, [])
check(4, [(([2], [3]), [4])])

def test_product(self):
for args, result in [
([], [()]), # zero iterables
Expand Down
Original file line number Diff line number Diff line change
@@ -0,0 +1,2 @@
Fix possible reference leaks and crash when re-enter the ``__next__()`` method of
:class:`itertools.pairwise`.
13 changes: 11 additions & 2 deletions Modules/itertoolsmodule.c
Original file line number Diff line number Diff line change
Expand Up @@ -330,21 +330,30 @@ pairwise_next(pairwiseobject *po)
return NULL;
}
if (old == NULL) {
po->old = old = (*Py_TYPE(it)->tp_iternext)(it);
old = (*Py_TYPE(it)->tp_iternext)(it);
Py_XSETREF(po->old, old);
if (old == NULL) {
Py_CLEAR(po->it);
return NULL;
}
it = po->it;
if (it == NULL) {
Py_CLEAR(po->old);
return NULL;
}
}
Py_INCREF(old);
new = (*Py_TYPE(it)->tp_iternext)(it);
if (new == NULL) {
Py_CLEAR(po->it);
Py_CLEAR(po->old);
Py_DECREF(old);
return NULL;
}
/* Future optimization: Reuse the result tuple as we do in enumerate() */
result = PyTuple_Pack(2, old, new);
Py_SETREF(po->old, new);
Py_XSETREF(po->old, new);
Py_DECREF(old);
return result;
}

Expand Down

0 comments on commit ea831c5

Please sign in to comment.