diff -r efade142ef01 Lib/test/test_pickle.py --- a/Lib/test/test_pickle.py Sun Jul 29 16:33:05 2012 +0200 +++ b/Lib/test/test_pickle.py Tue Jul 31 01:11:25 2012 +0300 @@ -1,6 +1,8 @@ import pickle import io +import struct +import unittest from test import support from test.pickletester import AbstractPickleTests @@ -101,6 +103,63 @@ pickler_class = _pickle.Pickler unpickler_class = _pickle.Unpickler + @support.cpython_only + class SizeofTests(unittest.TestCase): + check_sizeof = support.check_sizeof + + def test_pickler(self): + basesize = support.calcobjsize('5P2P2iP3iP') + p = _pickle.Pickler(io.BytesIO()) + self.assertEqual(object.__sizeof__(p), basesize) + MT_size = struct.calcsize('3PP') + ME_size = struct.calcsize('PP') + 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('2PP2P 2P2P2i4P2PP 2P3P6P2P2i') + unpickler = _pickle.Unpickler + P = struct.calcsize('P') + n = struct.calcsize('P') + 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] @@ -109,7 +168,7 @@ CDumpPickle_LoadPickle, DumpPickle_CLoadPickle, PyPicklerUnpicklerObjectTests, CPicklerUnpicklerObjectTests, - InMemoryPickleTests]) + InMemoryPickleTests, SizeofTests]) support.run_unittest(*tests) support.run_doctest(pickle) diff -r efade142ef01 Modules/_pickle.c --- a/Modules/_pickle.c Sun Jul 29 16:33:05 2012 +0200 +++ b/Modules/_pickle.c Tue Jul 31 01:11:25 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*/ }; @@ -3344,11 +3344,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 */ }; @@ -5480,11 +5498,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 */ };