diff --git a/Lib/test/test_itertools.py b/Lib/test/test_itertools.py --- a/Lib/test/test_itertools.py +++ b/Lib/test/test_itertools.py @@ -9,6 +9,7 @@ import random import copy import pickle +import struct from functools import reduce maxsize = support.MAX_Py_ssize_t minsize = -maxsize-1 @@ -1267,6 +1268,30 @@ self.pickletest(a, compare=ans) self.pickletest(b, compare=ans) + def test_tee_sizeof(self): + a, b = tee(range(10000)) + orig_size = sys.getsizeof(a) + self.assertEqual(sys.getsizeof(b), orig_size) + # When consuming a, a doesnt change size but b grows + take(1000, a) + self.assertEqual(sys.getsizeof(a), orig_size) + self.assertGreater(sys.getsizeof(b), orig_size) + size_b = sys.getsizeof(b) + # Again + take(1000, a) + self.assertEqual(sys.getsizeof(a), orig_size) + self.assertGreater(sys.getsizeof(b), size_b) + # Exhaust a + list(a) + self.assertEqual(sys.getsizeof(a), orig_size) + # 'b' holds all objects, so at least 10000 pointers + overhead + self.assertGreater(sys.getsizeof(b) - orig_size, + 10000 * struct.calcsize('P')) + # Exhausting b returns it to its original size, and a doesn't care + list(b) + self.assertEqual(sys.getsizeof(b), orig_size) + self.assertEqual(sys.getsizeof(a), orig_size) + # Issue 13454: Crash when deleting backward iterator from tee() def test_tee_del_backward(self): forward, backward = tee(repeat(None, 20000000)) diff --git a/Modules/itertoolsmodule.c b/Modules/itertoolsmodule.c --- a/Modules/itertoolsmodule.c +++ b/Modules/itertoolsmodule.c @@ -755,12 +755,29 @@ Py_RETURN_NONE; } +PyDoc_STRVAR(sizeof_doc, +"Return the memory size of the object in bytes"); + +static PyObject * +tee_sizeof(teeobject *to, void *unused) +{ + Py_ssize_t size = sizeof(teeobject); + teedataobject *tdo = to->dataobj; + while (tdo != NULL) { + assert(Py_TYPE(tdo) == &teedataobject_type); + size += sizeof(teedataobject); + tdo = (teedataobject *) tdo->nextlink; + } + return PyLong_FromSsize_t(size); +} + PyDoc_STRVAR(teeobject_doc, "Iterator wrapped to make it copyable"); static PyMethodDef tee_methods[] = { {"__copy__", (PyCFunction)tee_copy, METH_NOARGS, teecopy_doc}, {"__reduce__", (PyCFunction)tee_reduce, METH_NOARGS, reduce_doc}, + {"__sizeof__", (PyCFunction)tee_sizeof, METH_NOARGS, sizeof_doc}, {"__setstate__", (PyCFunction)tee_setstate, METH_O, setstate_doc}, {NULL, NULL} /* sentinel */ };