Skip to content

Commit ca6a36f

Browse files
hauntsaninjaSonicField
authored andcommitted
pythongh-118218: Reuse return tuple in itertools.pairwise (pythonGH-118219)
1 parent 692c94f commit ca6a36f

File tree

3 files changed

+40
-2
lines changed

3 files changed

+40
-2
lines changed

Lib/test/test_itertools.py

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1821,6 +1821,13 @@ def test_zip_longest_result_gc(self):
18211821
gc.collect()
18221822
self.assertTrue(gc.is_tracked(next(it)))
18231823

1824+
@support.cpython_only
1825+
def test_pairwise_result_gc(self):
1826+
# Ditto for pairwise.
1827+
it = pairwise([None, None])
1828+
gc.collect()
1829+
self.assertTrue(gc.is_tracked(next(it)))
1830+
18241831
@support.cpython_only
18251832
def test_immutable_types(self):
18261833
from itertools import _grouper, _tee, _tee_dataobject
Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1 @@
1+
Speed up :func:`itertools.pairwise` in the common case by up to 1.8x.

Modules/itertoolsmodule.c

Lines changed: 32 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -270,6 +270,7 @@ typedef struct {
270270
PyObject_HEAD
271271
PyObject *it;
272272
PyObject *old;
273+
PyObject *result;
273274
} pairwiseobject;
274275

275276
/*[clinic input]
@@ -301,6 +302,11 @@ pairwise_new_impl(PyTypeObject *type, PyObject *iterable)
301302
}
302303
po->it = it;
303304
po->old = NULL;
305+
po->result = PyTuple_Pack(2, Py_None, Py_None);
306+
if (po->result == NULL) {
307+
Py_DECREF(po);
308+
return NULL;
309+
}
304310
return (PyObject *)po;
305311
}
306312

@@ -311,6 +317,7 @@ pairwise_dealloc(pairwiseobject *po)
311317
PyObject_GC_UnTrack(po);
312318
Py_XDECREF(po->it);
313319
Py_XDECREF(po->old);
320+
Py_XDECREF(po->result);
314321
tp->tp_free(po);
315322
Py_DECREF(tp);
316323
}
@@ -321,6 +328,7 @@ pairwise_traverse(pairwiseobject *po, visitproc visit, void *arg)
321328
Py_VISIT(Py_TYPE(po));
322329
Py_VISIT(po->it);
323330
Py_VISIT(po->old);
331+
Py_VISIT(po->result);
324332
return 0;
325333
}
326334

@@ -355,8 +363,30 @@ pairwise_next(pairwiseobject *po)
355363
Py_DECREF(old);
356364
return NULL;
357365
}
358-
/* Future optimization: Reuse the result tuple as we do in enumerate() */
359-
result = PyTuple_Pack(2, old, new);
366+
367+
result = po->result;
368+
if (Py_REFCNT(result) == 1) {
369+
Py_INCREF(result);
370+
PyObject *last_old = PyTuple_GET_ITEM(result, 0);
371+
PyObject *last_new = PyTuple_GET_ITEM(result, 1);
372+
PyTuple_SET_ITEM(result, 0, Py_NewRef(old));
373+
PyTuple_SET_ITEM(result, 1, Py_NewRef(new));
374+
Py_DECREF(last_old);
375+
Py_DECREF(last_new);
376+
// bpo-42536: The GC may have untracked this result tuple. Since we're
377+
// recycling it, make sure it's tracked again:
378+
if (!_PyObject_GC_IS_TRACKED(result)) {
379+
_PyObject_GC_TRACK(result);
380+
}
381+
}
382+
else {
383+
result = PyTuple_New(2);
384+
if (result != NULL) {
385+
PyTuple_SET_ITEM(result, 0, Py_NewRef(old));
386+
PyTuple_SET_ITEM(result, 1, Py_NewRef(new));
387+
}
388+
}
389+
360390
Py_XSETREF(po->old, new);
361391
Py_DECREF(old);
362392
return result;

0 commit comments

Comments
 (0)