Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

gh-109786: Fix leaks when re-enter itertools.pairwise.__next__() #109788

Merged
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Jump to
Jump to file
Failed to load files.
Diff view
Diff view
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