diff -r d7a64e095930 Lib/test/test_itertools.py --- a/Lib/test/test_itertools.py Thu Jul 26 00:45:19 2012 +0200 +++ b/Lib/test/test_itertools.py Fri Jul 27 22:37:00 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 @@ -1453,6 +1456,68 @@ # we expect type errors because of wrong argument count self.assertNotIn("does not take keyword arguments", err.args[0]) +_header = '2P' +_vheader = _header + 'P' +_align = '0P' +if hasattr(sys, "gettotalrefcount"): + _header = '2P' + _header + _vheader = '2P' + _vheader + +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('P') + + 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('3PPi') + 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('3PPi') + 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('4PPi') + 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 @@ -1688,7 +1753,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 d7a64e095930 Modules/itertoolsmodule.c --- a/Modules/itertoolsmodule.c Thu Jul 26 00:45:19 2012 +0200 +++ b/Modules/itertoolsmodule.c Fri Jul 27 22:37:00 2012 +0300 @@ -1693,6 +1693,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) { @@ -1788,6 +1801,12 @@ return NULL; } +static PyMethodDef product_methods[] = { + {"__sizeof__", (PyCFunction)product_sizeof, METH_NOARGS, + sizeof_doc}, + {NULL, NULL} /* sentinel */ +}; + PyDoc_STRVAR(product_doc, "product(*iterables) --> product object\n\ \n\ @@ -1832,7 +1851,7 @@ 0, /* tp_weaklistoffset */ PyObject_SelfIter, /* tp_iter */ (iternextfunc)product_next, /* tp_iternext */ - 0, /* tp_methods */ + product_methods, /* tp_methods */ 0, /* tp_members */ 0, /* tp_getset */ 0, /* tp_base */ @@ -1925,6 +1944,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) { @@ -2019,6 +2049,12 @@ return NULL; } +static PyMethodDef combinations_methods[] = { + {"__sizeof__", (PyCFunction)combinations_sizeof, METH_NOARGS, + sizeof_doc}, + {NULL, NULL} /* sentinel */ +}; + PyDoc_STRVAR(combinations_doc, "combinations(iterable, r) --> combinations object\n\ \n\ @@ -2055,7 +2091,7 @@ 0, /* tp_weaklistoffset */ PyObject_SelfIter, /* tp_iter */ (iternextfunc)combinations_next, /* tp_iternext */ - 0, /* tp_methods */ + combinations_methods, /* tp_methods */ 0, /* tp_members */ 0, /* tp_getset */ 0, /* tp_base */ @@ -2174,6 +2210,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) { @@ -2264,6 +2311,12 @@ return NULL; } +static PyMethodDef cwr_methods[] = { + {"__sizeof__", (PyCFunction)cwr_sizeof, METH_NOARGS, + sizeof_doc}, + {NULL, NULL} /* sentinel */ +}; + PyDoc_STRVAR(cwr_doc, "combinations_with_replacement(iterable, r) --> combinations_with_replacement object\n\ \n\ @@ -2301,7 +2354,7 @@ 0, /* tp_weaklistoffset */ PyObject_SelfIter, /* tp_iter */ (iternextfunc)cwr_next, /* tp_iternext */ - 0, /* tp_methods */ + cwr_methods, /* tp_methods */ 0, /* tp_members */ 0, /* tp_getset */ 0, /* tp_base */ @@ -2437,6 +2490,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) { @@ -2536,6 +2602,12 @@ return NULL; } +static PyMethodDef permuations_methods[] = { + {"__sizeof__", (PyCFunction)permutations_sizeof, METH_NOARGS, + sizeof_doc}, + {NULL, NULL} /* sentinel */ +}; + PyDoc_STRVAR(permutations_doc, "permutations(iterable[, r]) --> permutations object\n\ \n\ @@ -2572,7 +2644,7 @@ 0, /* tp_weaklistoffset */ PyObject_SelfIter, /* tp_iter */ (iternextfunc)permutations_next, /* tp_iternext */ - 0, /* tp_methods */ + permuations_methods, /* tp_methods */ 0, /* tp_members */ 0, /* tp_getset */ 0, /* tp_base */