gh-92810: Reduce memory usage by ABCMeta.__subclasscheck__

Author: dolfinus

Lib/_py_abc.py

@@ -49,6 +49,14 @@ def __new__(mcls, name, bases, namespace, /, **kwargs):
         cls._abc_cache = WeakSet()
         cls._abc_negative_cache = WeakSet()
         cls._abc_negative_cache_version = ABCMeta._abc_invalidation_counter
+
+        # Performance optimization for common case
+        cls._abc_should_check_subclasses = False
+        if "__subclasses__" in namespace:
+            cls._abc_should_check_subclasses = True
+            for base in bases:
+                if hasattr(base, "_abc_should_check_subclasses"):
+                    base._abc_should_check_subclasses = True
         return cls
 
     def register(cls, subclass):
@@ -65,8 +73,20 @@ def register(cls, subclass):
         if issubclass(cls, subclass):
             # This would create a cycle, which is bad for the algorithm below
             raise RuntimeError("Refusing to create an inheritance cycle")
+        # Add registry entry
         cls._abc_registry.add(subclass)
         ABCMeta._abc_invalidation_counter += 1  # Invalidate negative cache
+        # Recursively register the subclass in all ABC bases,
+        # to avoid recursive lookups down the class tree.
+        # >>> class Ancestor1(ABC): pass
+        # >>> class Ancestor2(Ancestor1): pass
+        # >>> class Other: pass
+        # >>> Ancestor2.register(Other)  # calls Ancestor1.register(Other)
+        # >>> issubclass(Other, Ancestor2) is True
+        # >>> issubclass(Other, Ancestor1) is True  # already in registry
+        for base in cls.__bases__:
+            if hasattr(base, "_abc_registry"):
+                base.register(subclass)
         return subclass
 
     def _dump_registry(cls, file=None):
@@ -137,11 +157,16 @@ def __subclasscheck__(cls, subclass):
             if issubclass(subclass, rcls):
                 cls._abc_cache.add(subclass)
                 return True
-        # Check if it's a subclass of a subclass (recursive)
-        for scls in cls.__subclasses__():
-            if issubclass(subclass, scls):
-                cls._abc_cache.add(subclass)
-                return True
+        # Check if it's a subclass of a subclass (recursive).
+        # If __subclasses__ contain only ABCs,
+        # calling issubclass(...) will trigger the same __subclasscheck__
+        # on *every* element of class inheritance tree.
+        # Performing that only in resence of `def __subclasses__()` classmethod
+        if cls._abc_should_check_subclasses:
+            for scls in cls.__subclasses__():
+                if issubclass(subclass, scls):
+                    cls._abc_cache.add(subclass)
+                    return True
         # No dice; update negative cache
         cls._abc_negative_cache.add(subclass)
         return False

Lib/abc.py

@@ -104,7 +104,7 @@ class ABCMeta(type):
         """
         def __new__(mcls, name, bases, namespace, /, **kwargs):
             cls = super().__new__(mcls, name, bases, namespace, **kwargs)
-            _abc_init(cls)
+            _abc_init(cls, bases, namespace)
             return cls
 
         def register(cls, subclass):

Modules/_abc.c

@@ -57,6 +57,7 @@ typedef struct {
     PyObject *_abc_cache;
     PyObject *_abc_negative_cache;
     uint64_t _abc_negative_cache_version;
+    uint8_t _abc_should_check_subclasses;
 } _abc_data;
 
 #define _abc_data_CAST(op)  ((_abc_data *)(op))
@@ -73,6 +74,18 @@ set_cache_version(_abc_data *impl, uint64_t version)
     FT_ATOMIC_STORE_UINT64_RELAXED(impl->_abc_negative_cache_version, version);
 }
 
+static inline uint8_t
+get_should_check_subclasses(_abc_data *impl)
+{
+    return FT_ATOMIC_LOAD_UINT8_RELAXED(impl->_abc_should_check_subclasses);
+}
+
+static inline void
+set_should_check_subclasses(_abc_data *impl)
+{
+    FT_ATOMIC_STORE_UINT8_RELAXED(impl->_abc_should_check_subclasses, 1);
+}
+
 static int
 abc_data_traverse(PyObject *op, visitproc visit, void *arg)
 {
@@ -123,6 +136,7 @@ abc_data_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
     self->_abc_cache = NULL;
     self->_abc_negative_cache = NULL;
     self->_abc_negative_cache_version = get_invalidation_counter(state);
+    self->_abc_should_check_subclasses = 0;
     return (PyObject *) self;
 }
 
@@ -161,6 +175,30 @@ _get_impl(PyObject *module, PyObject *self)
     return (_abc_data *)impl;
 }
 
+/* If class is inherited from ABC, set data to point to internal ABC state of class, and return 1.
+   If object is not inherited from ABC, return 0.
+   If error is encountered, return -1.
+ */
+static int
+_get_optional_impl(_abcmodule_state *state, PyObject *self, _abc_data **data)
+{
+    assert(data != NULL);
+    PyObject *impl = NULL;
+    int res = PyObject_GetOptionalAttr(self, &_Py_ID(_abc_impl), &impl);
+    if (res <= 0) {
+        *data = NULL;
+        return res;
+    }
+    if (!Py_IS_TYPE(impl, state->_abc_data_type)) {
+        PyErr_SetString(PyExc_TypeError, "_abc_impl is set to a wrong type");
+        Py_DECREF(impl);
+        *data = NULL;
+        return -1;
+    }
+    *data = (_abc_data *)impl;
+    return 1;
+}
+
 static int
 _in_weak_set(_abc_data *impl, PyObject **pset, PyObject *obj)
 {
@@ -331,39 +369,34 @@ _abc__get_dump(PyObject *module, PyObject *self)
     }
     PyObject *res;
     Py_BEGIN_CRITICAL_SECTION(impl);
-    res = Py_BuildValue("NNNK",
+    res = Py_BuildValue("NNNKK",
                         PySet_New(impl->_abc_registry),
                         PySet_New(impl->_abc_cache),
                         PySet_New(impl->_abc_negative_cache),
-                        get_cache_version(impl));
+                        get_cache_version(impl),
+                        get_should_check_subclasses(impl));
     Py_END_CRITICAL_SECTION();
     Py_DECREF(impl);
     return res;
 }
 
 // Compute set of abstract method names.
 static int
-compute_abstract_methods(PyObject *self)
+compute_abstract_methods(PyObject *self, PyObject *bases, PyObject *ns)
 {
     int ret = -1;
     PyObject *abstracts = PyFrozenSet_New(NULL);
     if (abstracts == NULL) {
         return -1;
     }
 
-    PyObject *ns = NULL, *items = NULL, *bases = NULL;  // Py_XDECREF()ed on error.
-
     /* Stage 1: direct abstract methods. */
-    ns = PyObject_GetAttr(self, &_Py_ID(__dict__));
-    if (!ns) {
-        goto error;
-    }
-
     // We can't use PyDict_Next(ns) even when ns is dict because
     // _PyObject_IsAbstract() can mutate ns.
-    items = PyMapping_Items(ns);
+    PyObject *items = PyMapping_Items(ns);
     if (!items) {
-        goto error;
+        Py_DECREF(abstracts);
+        return -1;
     }
     assert(PyList_Check(items));
     for (Py_ssize_t pos = 0; pos < PyList_GET_SIZE(items); pos++) {
@@ -398,15 +431,6 @@ compute_abstract_methods(PyObject *self)
     }
 
     /* Stage 2: inherited abstract methods. */
-    bases = PyObject_GetAttr(self, &_Py_ID(__bases__));
-    if (!bases) {
-        goto error;
-    }
-    if (!PyTuple_Check(bases)) {
-        PyErr_SetString(PyExc_TypeError, "__bases__ is not tuple");
-        goto error;
-    }
-
     for (Py_ssize_t pos = 0; pos < PyTuple_GET_SIZE(bases); pos++) {
         PyObject *item = PyTuple_GET_ITEM(bases, pos);  // borrowed
         PyObject *base_abstracts, *iter;
@@ -459,31 +483,58 @@ compute_abstract_methods(PyObject *self)
     ret = 0;
 error:
     Py_DECREF(abstracts);
-    Py_XDECREF(ns);
-    Py_XDECREF(items);
-    Py_XDECREF(bases);
+    Py_DECREF(items);
     return ret;
 }
 
+/*
+ * Notify base classes that child one has __subclasses__ overriden.
+ * Used as performance optimization in __subclasscheck__
+ */
+static int
+_abc_notify_subclasses_override(_abcmodule_state *state, PyObject *data, PyObject *bases)
+{
+    set_should_check_subclasses((_abc_data*) data);
+
+    for (Py_ssize_t pos = 0; pos < PyTuple_GET_SIZE(bases); pos++) {
+        PyObject *base_class = PyTuple_GET_ITEM(bases, pos);  // borrowed
+        _abc_data *base_impl = NULL;
+        int base_is_abc = _get_optional_impl(state, base_class, &base_impl);
+        if (base_is_abc < 0) {
+            return -1;
+        }
+        if (base_is_abc == 0) {
+            continue;
+        }
+        set_should_check_subclasses(base_impl);
+        Py_DECREF(base_impl);
+    }
+
+    return 0;
+}
+
 #define COLLECTION_FLAGS (Py_TPFLAGS_SEQUENCE | Py_TPFLAGS_MAPPING)
 
 /*[clinic input]
 @permit_long_summary
 _abc._abc_init
 
     self: object
+    bases: object(subclass_of="&PyTuple_Type")
+    namespace: object(subclass_of="&PyDict_Type")
     /
 
 Internal ABC helper for class set-up. Should be never used outside abc module.
 [clinic start generated code]*/
 
 static PyObject *
-_abc__abc_init(PyObject *module, PyObject *self)
-/*[clinic end generated code: output=594757375714cda1 input=0b3513f947736d39]*/
+_abc__abc_init_impl(PyObject *module, PyObject *self, PyObject *bases,
+                    PyObject *namespace)
+/*[clinic end generated code: output=a410180fefc86056 input=a984e4f7d36d6298]*/
 {
     _abcmodule_state *state = get_abc_state(module);
     PyObject *data;
-    if (compute_abstract_methods(self) < 0) {
+    if (compute_abstract_methods(self, bases, namespace) < 0) {
         return NULL;
     }
 
@@ -492,6 +543,12 @@ _abc__abc_init(PyObject *module, PyObject *self)
     if (data == NULL) {
         return NULL;
     }
+    if (PyDict_ContainsString(namespace, "__subclasses__")) {
+        if (_abc_notify_subclasses_override(state, data, bases) < 0) {
+            Py_DECREF(data);
+            return NULL;
+        }
+    }
     if (PyObject_SetAttr(self, &_Py_ID(_abc_impl), data) < 0) {
         Py_DECREF(data);
         return NULL;
@@ -564,6 +621,7 @@ _abc__abc_register_impl(PyObject *module, PyObject *self, PyObject *subclass)
     if (result < 0) {
         return NULL;
     }
+    /* Add registry entry */
     _abc_data *impl = _get_impl(module, self);
     if (impl == NULL) {
         return NULL;
@@ -575,7 +633,43 @@ _abc__abc_register_impl(PyObject *module, PyObject *self, PyObject *subclass)
     Py_DECREF(impl);
 
     /* Invalidate negative cache */
-    increment_invalidation_counter(get_abc_state(module));
+    _abcmodule_state *state = get_abc_state(module);
+    increment_invalidation_counter(state);
+
+    /*
+     * Recursively register the subclass in all ABC bases,
+     * to avoid recursive lookups down the class tree.
+     * >>> class Ancestor1(ABC): pass
+     * >>> class Ancestor2(Ancestor1): pass
+     * >>> class Other: pass
+     * >>> Ancestor2.register(Other)  # calls Ancestor1.register(Other)
+     * >>> issubclass(Other, Ancestor2) is True
+     * >>> issubclass(Other, Ancestor1) is True  # already in registry
+     */
+    PyObject *bases = PyObject_GetAttr(self, &_Py_ID(__bases__));
+    if (!bases) {
+        return NULL;
+    }
+    if (!PyTuple_Check(bases)) {
+        PyErr_SetString(PyExc_TypeError, "__bases__ is not tuple");
+        goto error;
+    }
+    for (Py_ssize_t pos = 0; pos < PyTuple_GET_SIZE(bases); pos++) {
+        PyObject *base = PyTuple_GET_ITEM(bases, pos);  // borrowed
+        int base_is_abc = PyObject_HasAttrWithError(base, &_Py_ID(_abc_impl));
+        if (base_is_abc < 0) {
+            goto error;
+        }
+        if (base_is_abc == 0) {
+            continue;
+        }
+        PyObject *res = PyObject_CallMethod(base, "register", "O", subclass);
+        Py_XDECREF(res);
+        if (!res) {
+            goto error;
+        }
+    }
+    Py_DECREF(bases);
 
     /* Set Py_TPFLAGS_SEQUENCE or Py_TPFLAGS_MAPPING flag */
     if (PyType_Check(self)) {
@@ -588,6 +682,10 @@ _abc__abc_register_impl(PyObject *module, PyObject *self, PyObject *subclass)
         }
     }
     return Py_NewRef(subclass);
+
+error:
+    Py_DECREF(bases);
+    return NULL;
 }
 
 
@@ -788,31 +886,38 @@ _abc__abc_subclasscheck_impl(PyObject *module, PyObject *self,
         goto end;
     }
 
-    /* 6. Check if it's a subclass of a subclass (recursive). */
-    subclasses = PyObject_CallMethod(self, "__subclasses__", NULL);
-    if (subclasses == NULL) {
-        goto end;
-    }
-    if (!PyList_Check(subclasses)) {
-        PyErr_SetString(PyExc_TypeError, "__subclasses__() must return a list");
-        goto end;
-    }
-    for (pos = 0; pos < PyList_GET_SIZE(subclasses); pos++) {
-        PyObject *scls = PyList_GetItemRef(subclasses, pos);
-        if (scls == NULL) {
+    /* 6. Check if it's a subclass of a subclass (recursive).
+     * If __subclasses__ contain only ABCs,
+     * calling issubclass(...) will trigger the same __subclasscheck__
+     * on *every* element of class inheritance tree.
+     * Performing that only in resence of `def __subclasses__()` classmethod
+     */
+    if (get_should_check_subclasses(impl)) {
+        subclasses = PyObject_CallMethod(self, "__subclasses__", NULL);
+        if (subclasses == NULL) {
             goto end;
         }
-        int r = PyObject_IsSubclass(subclass, scls);
-        Py_DECREF(scls);
-        if (r > 0) {
-            if (_add_to_weak_set(impl, &impl->_abc_cache, subclass) < 0) {
-                goto end;
-            }
-            result = Py_True;
+        if (!PyList_Check(subclasses)) {
+            PyErr_SetString(PyExc_TypeError, "__subclasses__() must return a list");
             goto end;
         }
-        if (r < 0) {
-            goto end;
+        for (pos = 0; pos < PyList_GET_SIZE(subclasses); pos++) {
+            PyObject *scls = PyList_GetItemRef(subclasses, pos);
+            if (scls == NULL) {
+                goto end;
+            }
+            int r = PyObject_IsSubclass(subclass, scls);
+            Py_DECREF(scls);
+            if (r > 0) {
+                if (_add_to_weak_set(impl, &impl->_abc_cache, subclass) < 0) {
+                    goto end;
+                }
+                result = Py_True;
+                goto end;
+            }
+            if (r < 0) {
+                goto end;
+            }
         }
     }