diff -r fa95b04e67fd Lib/test/test_cpickle.py --- a/Lib/test/test_cpickle.py Sun Jul 29 16:30:50 2012 +0200 +++ b/Lib/test/test_cpickle.py Tue Jul 31 01:11:52 2012 +0300 @@ -1,8 +1,10 @@ import cPickle, unittest from cStringIO import StringIO +import io +import struct from test.pickletester import AbstractPickleTests, AbstractPickleModuleTests from test.pickletester import AbstractPicklerUnpicklerObjectTests -from test import test_support +from test import test_support as support class cPickleTests(AbstractPickleTests, AbstractPickleModuleTests): @@ -124,15 +126,48 @@ res.append(dict(doc=x, similar=[])) cPickle.dumps(res) +import sys +@support.cpython_only +class SizeofTests(unittest.TestCase): + check_sizeof = support.check_sizeof + + def test_pickler(self): + basesize = support.calcobjsize('7P3i2PiPiP') + pickler = cPickle.Pickler + check = self.check_sizeof + check(pickler(open(support.TESTFN, 'wb')), basesize) + check(pickler(None), basesize) + check(pickler(StringIO()), basesize) + WRITE_BUF_SIZE = 256 + check(pickler(io.BytesIO()), basesize + WRITE_BUF_SIZE) + + def test_unpickler(self): + basesize = support.calcobjsize('11P2i2Pi2P') + m_size = struct.calcsize('i') + + def check_unpickler(data, marks_size, buf_size): + with open(support.TESTFN, 'wb') as f: + cPickle.dump(data, f) + with open(support.TESTFN, 'rb') as f: + u = cPickle.Unpickler(f) + u.load() + self.check_sizeof(u, + basesize + marks_size * m_size + buf_size) + + check_unpickler(0, 0, 32) + check_unpickler([0] * 1000, 20, 32) + check_unpickler('a' * 100, 0, 128) + def test_main(): - test_support.run_unittest( + support.run_unittest( cPickleTests, cPicklePicklerTests, cPickleListPicklerTests, cPickleFastPicklerTests, cPickleDeepRecursive, cPicklePicklerUnpicklerObjectTests, + SizeofTests, ) if __name__ == "__main__": diff -r fa95b04e67fd Modules/cPickle.c --- a/Modules/cPickle.c Sun Jul 29 16:30:50 2012 +0200 +++ b/Modules/cPickle.c Tue Jul 31 01:11:52 2012 +0300 @@ -159,8 +159,11 @@ } static PyTypeObject PdataType = { - PyVarObject_HEAD_INIT(NULL, 0) "cPickle.Pdata", sizeof(Pdata), 0, - (destructor)Pdata_dealloc, + PyVarObject_HEAD_INIT(NULL, 0) + "_pickle.Pdata", /*tp_name*/ + sizeof(Pdata), /*tp_basicsize*/ + sizeof(PyObject *), /*tp_itemsize*/ + (destructor)Pdata_dealloc, /*tp_dealloc*/ 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0L,0L,0L,0L, "" }; @@ -2968,6 +2971,20 @@ return (PyObject*)self; } +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->write_buf != NULL) + res += WRITE_BUF_SIZE; + return PyLong_FromSsize_t(res); +} + static struct PyMethodDef Pickler_methods[] = { @@ -2978,6 +2995,8 @@ PyDoc_STR("clear_memo() -- Clear the picklers memo")}, {"getvalue", (PyCFunction)Pickle_getvalue, METH_VARARGS, PyDoc_STR("getvalue() -- Finish picking a list-based pickle")}, + {"__sizeof__", (PyCFunction)Pickler_sizeof, METH_NOARGS, + sizeof_doc}, {NULL, NULL} /* sentinel */ }; @@ -5399,6 +5418,18 @@ return noload(self); } +static PyObject * +Unpickler_sizeof(Unpicklerobject *self, void *unused) +{ + Py_ssize_t res; + + res = sizeof(Unpicklerobject); + if (self->marks != NULL) + res += self->marks_size * sizeof(Py_ssize_t); + res += self->buf_size; + return PyLong_FromSsize_t(res); +} + static struct PyMethodDef Unpickler_methods[] = { {"load", (PyCFunction)Unpickler_load, METH_NOARGS, @@ -5413,6 +5444,8 @@ "persistent references without instantiating any objects or importing\n" "any modules.\n") }, + {"__sizeof__", (PyCFunction)Unpickler_sizeof, METH_NOARGS, + sizeof_doc}, {NULL, NULL} /* sentinel */ }; @@ -5769,6 +5802,12 @@ Unpicklertype__doc__, /* tp_doc */ (traverseproc)Unpickler_traverse, /* tp_traverse */ (inquiry)Unpickler_clear, /* tp_clear */ + + 0, /* tp_richcompare */ + 0, /* tp_weaklistoffset */ + 0, /* tp_iter */ + 0, /* tp_iternext */ + Unpickler_methods, /* tp_methods */ }; static struct PyMethodDef cPickle_methods[] = {