diff --git a/Lib/test/test_itertools.py b/Lib/test/test_itertools.py index 95e67911db6..e243da309f0 100644 --- a/Lib/test/test_itertools.py +++ b/Lib/test/test_itertools.py @@ -1821,6 +1821,13 @@ class TestBasicOps(unittest.TestCase): gc.collect() self.assertTrue(gc.is_tracked(next(it))) + @support.cpython_only + def test_pairwise_result_gc(self): + # Ditto for pairwise. + it = pairwise([None, None]) + gc.collect() + self.assertTrue(gc.is_tracked(next(it))) + @support.cpython_only def test_immutable_types(self): from itertools import _grouper, _tee, _tee_dataobject diff --git a/Misc/NEWS.d/next/Library/2024-04-24-07-45-08.gh-issue-118218.m1OHbN.rst b/Misc/NEWS.d/next/Library/2024-04-24-07-45-08.gh-issue-118218.m1OHbN.rst new file mode 100644 index 00000000000..6d3c54cb485 --- /dev/null +++ b/Misc/NEWS.d/next/Library/2024-04-24-07-45-08.gh-issue-118218.m1OHbN.rst @@ -0,0 +1 @@ +Speed up :func:`itertools.pairwise` in the common case by up to 1.8x. diff --git a/Modules/itertoolsmodule.c b/Modules/itertoolsmodule.c index 6ee447ef6a8..21ce3ecfad0 100644 --- a/Modules/itertoolsmodule.c +++ b/Modules/itertoolsmodule.c @@ -270,6 +270,7 @@ typedef struct { PyObject_HEAD PyObject *it; PyObject *old; + PyObject *result; } pairwiseobject; /*[clinic input] @@ -301,6 +302,11 @@ pairwise_new_impl(PyTypeObject *type, PyObject *iterable) } po->it = it; po->old = NULL; + po->result = PyTuple_Pack(2, Py_None, Py_None); + if (po->result == NULL) { + Py_DECREF(po); + return NULL; + } return (PyObject *)po; } @@ -311,6 +317,7 @@ pairwise_dealloc(pairwiseobject *po) PyObject_GC_UnTrack(po); Py_XDECREF(po->it); Py_XDECREF(po->old); + Py_XDECREF(po->result); tp->tp_free(po); Py_DECREF(tp); } @@ -321,6 +328,7 @@ pairwise_traverse(pairwiseobject *po, visitproc visit, void *arg) Py_VISIT(Py_TYPE(po)); Py_VISIT(po->it); Py_VISIT(po->old); + Py_VISIT(po->result); return 0; } @@ -355,8 +363,30 @@ pairwise_next(pairwiseobject *po) Py_DECREF(old); return NULL; } - /* Future optimization: Reuse the result tuple as we do in enumerate() */ - result = PyTuple_Pack(2, old, new); + + result = po->result; + if (Py_REFCNT(result) == 1) { + Py_INCREF(result); + PyObject *last_old = PyTuple_GET_ITEM(result, 0); + PyObject *last_new = PyTuple_GET_ITEM(result, 1); + PyTuple_SET_ITEM(result, 0, Py_NewRef(old)); + PyTuple_SET_ITEM(result, 1, Py_NewRef(new)); + Py_DECREF(last_old); + Py_DECREF(last_new); + // bpo-42536: The GC may have untracked this result tuple. Since we're + // recycling it, make sure it's tracked again: + if (!_PyObject_GC_IS_TRACKED(result)) { + _PyObject_GC_TRACK(result); + } + } + else { + result = PyTuple_New(2); + if (result != NULL) { + PyTuple_SET_ITEM(result, 0, Py_NewRef(old)); + PyTuple_SET_ITEM(result, 1, Py_NewRef(new)); + } + } + Py_XSETREF(po->old, new); Py_DECREF(old); return result;