Account for shared keys in type's __sizeof__ (#13903).

loewis · loewis · commit 4f2f3b6217a5 · 2012-04-24T19:13:57.000+02:00
diff --git a/Include/dictobject.h b/Include/dictobject.h
@@ -75,6 +75,7 @@ PyAPI_FUNC(int) _PyDict_Contains(PyObject *mp, PyObject *key, Py_hash_t hash);
 PyAPI_FUNC(PyObject *) _PyDict_NewPresized(Py_ssize_t minused);
 PyAPI_FUNC(void) _PyDict_MaybeUntrack(PyObject *mp);
 PyAPI_FUNC(int) _PyDict_HasOnlyStringKeys(PyObject *mp);
+Py_ssize_t _PyDict_KeysSize(PyDictKeysObject *keys);
 #define _PyDict_HasSplitTable(d) ((d)->ma_values != NULL)
 
 PyAPI_FUNC(int) PyDict_ClearFreeList(void);
diff --git a/Lib/test/test_sys.py b/Lib/test/test_sys.py
@@ -829,13 +829,19 @@ def delx(self): del self.__x
         check((), size(vh))
         check((1,2,3), size(vh) + 3*self.P)
         # type
-        # (PyTypeObject + PyNumberMethods + PyMappingMethods +
-        #  PySequenceMethods + PyBufferProcs)
-        s = size(vh + 'P2P15Pl4PP9PP11PIP') + size('16Pi17P 3P 10P 2P 3P')
+        # static type: PyTypeObject
+        s = size(vh + 'P2P15Pl4PP9PP11PI')
         check(int, s)
+        # (PyTypeObject + PyNumberMethods + PyMappingMethods +
+        #  PySequenceMethods + PyBufferProcs + 4P)
+        s = size(vh + 'P2P15Pl4PP9PP11PI') + size('34P 3P 10P 2P 4P')
+        # Separate block for PyDictKeysObject with 4 entries
+        s += size("PPPP") + 4*size("PPP")
         # class
         class newstyleclass(object): pass
         check(newstyleclass, s)
+        # dict with shared keys
+        check(newstyleclass().__dict__, size(h+"PPP4P"))
         # unicode
         # each tuple contains a string and its expected character size
         # don't put any static strings here, as they may contain
diff --git a/Objects/dictobject.c b/Objects/dictobject.c
@@ -2398,22 +2398,23 @@ static PyObject *dictiter_new(PyDictObject *, PyTypeObject *);
 static PyObject *
 dict_sizeof(PyDictObject *mp)
 {
-    Py_ssize_t size;
-    double res, keys_size;
+    Py_ssize_t size, res;
 
     size = DK_SIZE(mp->ma_keys);
     res = sizeof(PyDictObject);
     if (mp->ma_values)
         res += size * sizeof(PyObject*);
-    /* Count our share of the keys object -- with rounding errors. */
-    keys_size = sizeof(PyDictKeysObject) + (size-1) * sizeof(PyDictKeyEntry);
-    /* If refcnt > 1, then one count is (probably) held by a type */
-    /* XXX  This is somewhat approximate :) */
-    if (mp->ma_keys->dk_refcnt < 3)
-        res += keys_size;
-    else
-        res += keys_size / (mp->ma_keys->dk_refcnt - 1);
-    return PyFloat_FromDouble(res);
+    /* If the dictionary is split, the keys portion is accounted-for
+       in the type object. */
+    if (mp->ma_keys->dk_refcnt == 1)
+        res += sizeof(PyDictKeysObject) + (size-1) * sizeof(PyDictKeyEntry);
+    return PyLong_FromSsize_t(res);
+}
+
+Py_ssize_t
+_PyDict_KeysSize(PyDictKeysObject *keys)
+{
+    return sizeof(PyDictKeysObject) + (DK_SIZE(keys)-1) * sizeof(PyDictKeyEntry);
 }
 
 PyDoc_STRVAR(contains__doc__,
diff --git a/Objects/typeobject.c b/Objects/typeobject.c
@@ -2730,6 +2730,22 @@ type_dir(PyObject *self, PyObject *args)
     return result;
 }
 
+static PyObject*
+type_sizeof(PyObject *self, PyObject *args_unused)
+{
+    Py_ssize_t size;
+    PyTypeObject *type = (PyTypeObject*)self;
+    if (type->tp_flags & Py_TPFLAGS_HEAPTYPE) {
+        PyHeapTypeObject* et = (PyHeapTypeObject*)type;
+        size = sizeof(PyHeapTypeObject);
+        if (et->ht_cached_keys)
+            size += _PyDict_KeysSize(et->ht_cached_keys);
+    }
+    else
+        size = sizeof(PyTypeObject);
+    return PyLong_FromSsize_t(size);
+}
+
 static PyMethodDef type_methods[] = {
     {"mro", (PyCFunction)mro_external, METH_NOARGS,
      PyDoc_STR("mro() -> list\nreturn a type's method resolution order")},
@@ -2745,6 +2761,8 @@ static PyMethodDef type_methods[] = {
      PyDoc_STR("__subclasscheck__() -> bool\ncheck if a class is a subclass")},
     {"__dir__", type_dir, METH_NOARGS,
      PyDoc_STR("__dir__() -> list\nspecialized __dir__ implementation for types")},
+    {"__sizeof__",      type_sizeof,       METH_NOARGS,
+     "__sizeof__() -> int\nreturn memory consumption of the type object"},
     {0}
 };
 
