gh-118218: Reuse return tuple in itertools.pairwise (GH-118219)

This commit is contained in:
Shantanu 2024-04-30 13:16:52 -07:00 committed by GitHub
parent b568c2c1ff
commit 6999d68d28
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
3 changed files with 40 additions and 2 deletions

View File

@ -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

View File

@ -0,0 +1 @@
Speed up :func:`itertools.pairwise` in the common case by up to 1.8x.

View File

@ -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;