diff -r e77c117ab024 Lib/test/test_pickle.py --- a/Lib/test/test_pickle.py Mon Jul 30 13:58:42 2012 +0200 +++ b/Lib/test/test_pickle.py Tue Jul 31 01:11:16 2012 +0300 @@ -1,7 +1,9 @@ import pickle import io import collections +import struct +import unittest from test import support from test.pickletester import AbstractPickleTests @@ -125,6 +127,63 @@ def get_dispatch_table(self): return collections.ChainMap({}, pickle.dispatch_table) + @support.cpython_only + class SizeofTests(unittest.TestCase): + check_sizeof = support.check_sizeof + + def test_pickler(self): + basesize = support.calcobjsize('6P2n2in3iP') + p = _pickle.Pickler(io.BytesIO()) + self.assertEqual(object.__sizeof__(p), basesize) + MT_size = struct.calcsize('3nP0n') + ME_size = struct.calcsize('Pn0P') + check = self.check_sizeof + check(p, basesize + MT_size + 8 * ME_size) + for i in range(0, 256): + p.dump(chr(i)) + check(p, basesize + MT_size + 512 * ME_size) + + def test_unpickler(self): + basesize = support.calcobjsize('2Pn2P 2P2n2i5P 2P3n6P2n2i') + unpickler = _pickle.Unpickler + P = struct.calcsize('P') + n = struct.calcsize('n') + check = self.check_sizeof + for encoding in 'ASCII', 'UTF-16', 'latin-1': + for errors in 'strict', 'replace': + u = unpickler(io.BytesIO(), + encoding=encoding, errors=errors) + self.assertEqual(object.__sizeof__(u), basesize) + check(u, basesize + 32 * P + + len(encoding) + 1 + len(errors) + 1) + + stdsize = basesize + len('ASCII') + 1 + len('strict') + 1 + def check_unpickler(data, memo_size, marks_size): + dump = pickle.dumps(data) + u = unpickler(io.BytesIO(dump), + encoding='ASCII', errors='strict') + u.load() + check(u, stdsize + memo_size * P + marks_size * n) + + check_unpickler(0, 32, 0) + check_unpickler([0] * 1000, 32, 20) + check_unpickler([chr(i) for i in range(1000)], 1024, 20) + def recurse(deep): + data = 0 + for i in range(deep): + data = [data, data] + return data + check_unpickler(recurse(0), 32, 0) + check_unpickler(recurse(1), 32, 20) + check_unpickler(recurse(20), 32, 58) + check_unpickler(recurse(50), 64, 58) + check_unpickler(recurse(100), 128, 134) + + u = unpickler(io.BytesIO(pickle.dumps('a', 0)), + encoding='ASCII', errors='strict') + u.load() + check(u, stdsize + 32 * P + 2 + 1) + def test_main(): tests = [PickleTests, PyPicklerTests, PyPersPicklerTests, @@ -135,7 +194,7 @@ PyPicklerUnpicklerObjectTests, CPicklerUnpicklerObjectTests, CDispatchTableTests, CChainDispatchTableTests, - InMemoryPickleTests]) + InMemoryPickleTests, SizeofTests]) support.run_unittest(*tests) support.run_doctest(pickle) diff -r e77c117ab024 Modules/_pickle.c --- a/Modules/_pickle.c Mon Jul 30 13:58:42 2012 +0200 +++ b/Modules/_pickle.c Tue Jul 31 01:11:16 2012 +0300 @@ -165,7 +165,7 @@ PyVarObject_HEAD_INIT(NULL, 0) "_pickle.Pdata", /*tp_name*/ sizeof(Pdata), /*tp_basicsize*/ - 0, /*tp_itemsize*/ + sizeof(PyObject *), /*tp_itemsize*/ (destructor)Pdata_dealloc, /*tp_dealloc*/ }; @@ -3384,11 +3384,29 @@ Py_RETURN_NONE; } +PyDoc_STRVAR(sizeof_doc, +"__sizeof__() -> int. Returns size in memory, in bytes."); + +static PyObject * +Pickler_sizeof(PicklerObject *self, void *unused) +{ + Py_ssize_t res; + + res = sizeof(PicklerObject); + if (self->memo != NULL) { + res += sizeof(PyMemoTable); + res += self->memo->mt_allocated * sizeof(PyMemoEntry); + } + return PyLong_FromSsize_t(res); +} + static struct PyMethodDef Pickler_methods[] = { {"dump", (PyCFunction)Pickler_dump, METH_VARARGS, Pickler_dump_doc}, {"clear_memo", (PyCFunction)Pickler_clear_memo, METH_NOARGS, Pickler_clear_memo_doc}, + {"__sizeof__", (PyCFunction)Pickler_sizeof, METH_NOARGS, + sizeof_doc}, {NULL, NULL} /* sentinel */ }; @@ -5541,11 +5559,32 @@ return global; } +static PyObject * +Unpickler_sizeof(UnpicklerObject *self, void *unused) +{ + Py_ssize_t res; + + res = sizeof(UnpicklerObject); + if (self->memo != NULL) + res += self->memo_size * sizeof(PyObject *); + if (self->marks != NULL) + res += self->marks_size * sizeof(Py_ssize_t); + if (self->input_line != NULL) + res += strlen(self->input_line) + 1; + if (self->encoding != NULL) + res += strlen(self->encoding) + 1; + if (self->errors != NULL) + res += strlen(self->errors) + 1; + return PyLong_FromSsize_t(res); +} + static struct PyMethodDef Unpickler_methods[] = { {"load", (PyCFunction)Unpickler_load, METH_NOARGS, Unpickler_load_doc}, {"find_class", (PyCFunction)Unpickler_find_class, METH_VARARGS, Unpickler_find_class_doc}, + {"__sizeof__", (PyCFunction)Unpickler_sizeof, METH_NOARGS, + sizeof_doc}, {NULL, NULL} /* sentinel */ };