Skip to content

Commit

Permalink
pythongh-109786: Fix leaks when re-enter itertools.pairwise.__next__()
Browse files Browse the repository at this point in the history
  • Loading branch information
serhiy-storchaka committed Sep 23, 2023
1 parent e8be0c9 commit c0f4360
Show file tree
Hide file tree
Showing 3 changed files with 42 additions and 2 deletions.
32 changes: 32 additions & 0 deletions Lib/test/test_itertools.py
Expand Up @@ -1152,6 +1152,38 @@ 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 == 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], ([3], [4])),
(([3], [4]), [5]),
([5], [6]),
])
check(3, [
([1], [2]),
([2], ([2], [4])),
(([2], [4]), [5]),
([5], [6]),
])

def test_product(self):
for args, result in [
([], [()]), # zero iterables
Expand Down
@@ -0,0 +1,2 @@
Fix reference leaks when re-enter the ``__next__()`` method of
:class:`itertools.pairwise`.
10 changes: 8 additions & 2 deletions Modules/itertoolsmodule.c
Expand Up @@ -330,21 +330,27 @@ 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);
if (old == NULL) {
Py_CLEAR(po->it);
Py_CLEAR(po->old);
return NULL;
}
}
else {
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 c0f4360

Please sign in to comment.