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

2 files changed, 30 insertions, 11 deletions

diff --git a/Objects/dictobject.c b/Objects/dictobject.c
index 6eb9b25..610a5ee 100644
--- 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
index fb48ce2..58f55df 100644
--- 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}
 };