diff -r f691b32cc1fd Lib/test/test_itertools.py --- a/Lib/test/test_itertools.py Fri Jul 27 14:55:17 2012 +0100 +++ b/Lib/test/test_itertools.py Fri Jul 27 22:39:55 2012 +0300 @@ -10,6 +10,9 @@ import copy import pickle from functools import reduce +import sys +import struct +import _testcapi maxsize = support.MAX_Py_ssize_t minsize = -maxsize-1 @@ -1802,6 +1805,69 @@ # we expect type errors because of wrong argument count self.assertNotIn("does not take keyword arguments", err.args[0]) +_header = 'nP' +_vheader = _header + 'n' +_align = '0n' +if hasattr(sys, "gettotalrefcount"): + _header = '2P' + _header + _vheader = '2P' + _vheader + _align = '0P' + +def calcobjsize(fmt): + return struct.calcsize(_header + fmt + _align) + +def calcvobjsize(fmt): + return struct.calcsize(_vheader + fmt + _align) + +_TPFLAGS_HAVE_GC = 1<<14 +_TPFLAGS_HEAPTYPE = 1<<9 + +@support.cpython_only +class SizeofTest(unittest.TestCase): + def setUp(self): + self.ssize_t = struct.calcsize('n') + + def check_sizeof(self, o, size): + result = sys.getsizeof(o) + # add GC header size + if ((type(o) == type) and (o.__flags__ & _TPFLAGS_HEAPTYPE) or\ + ((type(o) != type) and (type(o).__flags__ & _TPFLAGS_HAVE_GC))): + size += _testcapi.SIZEOF_PYGC_HEAD + msg = 'wrong size for %s: got %d, expected %d' \ + % (type(o), result, size) + self.assertEqual(result, size, msg) + + def test_product_sizeof(self): + basesize = calcobjsize('3Pi') + check = self.check_sizeof + check(product('ab', '12'), basesize + 2 * self.ssize_t) + check(product(*(('abc',) * 10)), basesize + 10 * self.ssize_t) + + def test_combinations_sizeof(self): + basesize = calcobjsize('3Pni') + check = self.check_sizeof + check(combinations('abcd', 3), basesize + 3 * self.ssize_t) + check(combinations(range(10), 4), basesize + 4 * self.ssize_t) + + def test_combinations_with_replacement_sizeof(self): + cwr = combinations_with_replacement + basesize = calcobjsize('3Pni') + check = self.check_sizeof + check(cwr('abcd', 3), basesize + 3 * self.ssize_t) + check(cwr(range(10), 4), basesize + 4 * self.ssize_t) + + def test_permutations_sizeof(self): + basesize = calcobjsize('4Pni') + check = self.check_sizeof + check(permutations('abcd'), + basesize + 4 * self.ssize_t + 4 * self.ssize_t) + check(permutations('abcd', 3), + basesize + 4 * self.ssize_t + 3 * self.ssize_t) + check(permutations('abcde', 3), + basesize + 5 * self.ssize_t + 3 * self.ssize_t) + check(permutations(range(10), 4), + basesize + 10 * self.ssize_t + 4 * self.ssize_t) + libreftest = """ Doctest for examples in the library reference: libitertools.tex @@ -2037,7 +2103,8 @@ def test_main(verbose=None): test_classes = (TestBasicOps, TestVariousIteratorArgs, TestGC, RegressionTests, LengthTransparency, - SubclassWithKwargsTest, TestExamples) + SubclassWithKwargsTest, TestExamples, + SizeofTest) support.run_unittest(*test_classes) # verify reference counting diff -r f691b32cc1fd Modules/itertoolsmodule.c --- a/Modules/itertoolsmodule.c Fri Jul 27 14:55:17 2012 +0100 +++ b/Modules/itertoolsmodule.c Fri Jul 27 22:39:55 2012 +0300 @@ -945,7 +945,7 @@ /* Create a new cycle with the iterator tuple, then set * the saved state on it. */ - return Py_BuildValue("O(O)(Oi)", Py_TYPE(lz), + return Py_BuildValue("O(O)(Oi)", Py_TYPE(lz), lz->it, lz->saved, lz->firstpass); } @@ -2034,6 +2034,19 @@ Py_TYPE(lz)->tp_free(lz); } +static PyObject * +product_sizeof(productobject *lz, void *unused) +{ + Py_ssize_t res; + + res = sizeof(productobject); + if (lz->indices != NULL && lz->pools != NULL) + res += PyTuple_GET_SIZE(lz->pools) * sizeof(Py_ssize_t); + return PyLong_FromSsize_t(res); +} + +PyDoc_STRVAR(sizeof_doc, "Returns size in memory, in bytes."); + static int product_traverse(productobject *lz, visitproc visit, void *arg) { @@ -2203,6 +2216,8 @@ reduce_doc}, {"__setstate__", (PyCFunction)product_setstate, METH_O, setstate_doc}, + {"__sizeof__", (PyCFunction)product_sizeof, METH_NOARGS, + sizeof_doc}, {NULL, NULL} /* sentinel */ }; @@ -2343,6 +2358,17 @@ Py_TYPE(co)->tp_free(co); } +static PyObject * +combinations_sizeof(combinationsobject *co, void *unused) +{ + Py_ssize_t res; + + res = sizeof(combinationsobject); + if (co->indices != NULL) + res += co->r * sizeof(Py_ssize_t); + return PyLong_FromSsize_t(res); +} + static int combinations_traverse(combinationsobject *co, visitproc visit, void *arg) { @@ -2514,6 +2540,8 @@ reduce_doc}, {"__setstate__", (PyCFunction)combinations_setstate, METH_O, setstate_doc}, + {"__sizeof__", (PyCFunction)combinations_sizeof, METH_NOARGS, + sizeof_doc}, {NULL, NULL} /* sentinel */ }; @@ -2672,6 +2700,17 @@ Py_TYPE(co)->tp_free(co); } +static PyObject * +cwr_sizeof(cwrobject *co, void *unused) +{ + Py_ssize_t res; + + res = sizeof(cwrobject); + if (co->indices != NULL) + res += co->r * sizeof(Py_ssize_t); + return PyLong_FromSsize_t(res); +} + static int cwr_traverse(cwrobject *co, visitproc visit, void *arg) { @@ -2835,6 +2874,8 @@ reduce_doc}, {"__setstate__", (PyCFunction)cwr_setstate, METH_O, setstate_doc}, + {"__sizeof__", (PyCFunction)cwr_sizeof, METH_NOARGS, + sizeof_doc}, {NULL, NULL} /* sentinel */ }; @@ -3011,6 +3052,19 @@ Py_TYPE(po)->tp_free(po); } +static PyObject * +permutations_sizeof(permutationsobject *po, void *unused) +{ + Py_ssize_t res; + + res = sizeof(permutationsobject); + if (po->indices != NULL && po->pool != NULL) + res += PyTuple_GET_SIZE(po->pool) * sizeof(Py_ssize_t); + if (po->cycles != NULL) + res += po->r * sizeof(Py_ssize_t); + return PyLong_FromSsize_t(res); +} + static int permutations_traverse(permutationsobject *po, visitproc visit, void *arg) { @@ -3132,7 +3186,7 @@ goto err; PyTuple_SET_ITEM(indices, i, index); } - + cycles = PyTuple_New(po->r); if (cycles == NULL) goto err; @@ -3158,7 +3212,7 @@ { PyObject *indices, *cycles, *result; Py_ssize_t n, i; - + if (!PyArg_ParseTuple(state, "O!O!", &PyTuple_Type, &indices, &PyTuple_Type, &cycles)) @@ -3216,6 +3270,8 @@ reduce_doc}, {"__setstate__", (PyCFunction)permutations_setstate, METH_O, setstate_doc}, + {"__sizeof__", (PyCFunction)permutations_sizeof, METH_NOARGS, + sizeof_doc}, {NULL, NULL} /* sentinel */ }; @@ -3337,18 +3393,18 @@ accumulate_next(accumulateobject *lz) { PyObject *val, *oldtotal, *newtotal; - + val = PyIter_Next(lz->it); if (val == NULL) return NULL; - + if (lz->total == NULL) { Py_INCREF(val); lz->total = val; return lz->total; } - if (lz->binop == NULL) + if (lz->binop == NULL) newtotal = PyNumber_Add(lz->total, val); else newtotal = PyObject_CallFunctionObjArgs(lz->binop, lz->total, val, NULL); @@ -3359,7 +3415,7 @@ oldtotal = lz->total; lz->total = newtotal; Py_DECREF(oldtotal); - + Py_INCREF(newtotal); return newtotal; } @@ -4327,7 +4383,7 @@ static PyObject * zip_longest_reduce(ziplongestobject *lz) { - + /* Create a new tuple with empty sequences where appropriate to pickle. * Then use setstate to set the fillvalue */