gh-92810: Avoid O(n^2) complexity in ABCMeta.__subclasscheck__

Signed-off-by: Martynov Maxim <[email protected]>

Author: dolfinus

Lib/_py_abc.py

@@ -1,5 +1,7 @@
 from _weakrefset import WeakSet
 
+_UNSET = object()
+
 
 def get_cache_token():
     """Returns the current ABC cache token.
@@ -65,8 +67,23 @@ 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")
+
+        # Actual registration
         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.
+        # >>> class Ancestor1(ABC): pass
+        # >>> class Ancestor2(Ancestor1): pass
+        # >>> class Other: pass
+        # >>> Ancestor2.register(Other)  # same result for Ancestor1.register(Other)
+        # >>> issubclass(Other, Ancestor2) is True
+        # >>> issubclass(Other, Ancestor1) is True
+        for pcls in cls.__mro__:
+            if hasattr(pcls, "_abc_registry"):
+                pcls._abc_registry.add(subclass)
+
+        # Invalidate negative cache
+        ABCMeta._abc_invalidation_counter += 1
         return subclass
 
     def _dump_registry(cls, file=None):
@@ -137,11 +154,19 @@ 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).
+        # >>> class Ancestor: __subclasses__ = lambda: [Other]
+        # >>> class Other: pass
+        # >>> isinstance(Other, Ancestor) is True
+        # Do not iterate over cls.__subclasses__() because it returns the entire class tree,
+        # not just direct children, which leads to O(n^2) lookup.
+        original_subclasses = getattr(cls, "__dict__", {}).get("__subclasses__", _UNSET)
+        if original_subclasses is not _UNSET:
+            for scls in original_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/test/test_abc.py

@@ -411,7 +411,35 @@ class MyInt(int):
             self.assertIsInstance(42, A)
             self.assertIsInstance(42, (A,))
 
-        def test_issubclass_bad_arguments(self):
+        def test_subclasses(self):
+            class A:
+                pass
+
+            class B:
+                pass
+
+            class C:
+                pass
+
+            class Sup(metaclass=abc_ABCMeta):
+                __subclasses__ = lambda: [A, B]
+
+            self.assertIsSubclass(A, Sup)
+            self.assertIsSubclass(A, (Sup,))
+            self.assertIsInstance(A(), Sup)
+            self.assertIsInstance(A(), (Sup,))
+
+            self.assertIsSubclass(B, Sup)
+            self.assertIsSubclass(B, (Sup,))
+            self.assertIsInstance(B(), Sup)
+            self.assertIsInstance(B(), (Sup,))
+
+            self.assertNotIsSubclass(C, Sup)
+            self.assertNotIsSubclass(C, (Sup,))
+            self.assertNotIsInstance(C(), Sup)
+            self.assertNotIsInstance(C(), (Sup,))
+
+        def test_subclasses_bad_arguments(self):
             class A(metaclass=abc_ABCMeta):
                 pass
 

Lib/test/test_abstract_numbers.py

@@ -24,11 +24,11 @@ def not_implemented(*args, **kwargs):
 
 class TestNumbers(unittest.TestCase):
     def test_int(self):
-        self.assertTrue(issubclass(int, Integral))
-        self.assertTrue(issubclass(int, Rational))
-        self.assertTrue(issubclass(int, Real))
-        self.assertTrue(issubclass(int, Complex))
-        self.assertTrue(issubclass(int, Number))
+        self.assertIsSubclass(int, Integral)
+        self.assertIsSubclass(int, Rational)
+        self.assertIsSubclass(int, Real)
+        self.assertIsSubclass(int, Complex)
+        self.assertIsSubclass(int, Number)
 
         self.assertEqual(7, int(7).real)
         self.assertEqual(0, int(7).imag)
@@ -38,23 +38,23 @@ def test_int(self):
         self.assertEqual(1, int(7).denominator)
 
     def test_float(self):
-        self.assertFalse(issubclass(float, Integral))
-        self.assertFalse(issubclass(float, Rational))
-        self.assertTrue(issubclass(float, Real))
-        self.assertTrue(issubclass(float, Complex))
-        self.assertTrue(issubclass(float, Number))
+        self.assertNotIsSubclass(float, Integral)
+        self.assertNotIsSubclass(float, Rational)
+        self.assertIsSubclass(float, Real)
+        self.assertIsSubclass(float, Complex)
+        self.assertIsSubclass(float, Number)
 
         self.assertEqual(7.3, float(7.3).real)
         self.assertEqual(0, float(7.3).imag)
         self.assertEqual(7.3, float(7.3).conjugate())
         self.assertEqual(-7.3, float(-7.3).conjugate())
 
     def test_complex(self):
-        self.assertFalse(issubclass(complex, Integral))
-        self.assertFalse(issubclass(complex, Rational))
-        self.assertFalse(issubclass(complex, Real))
-        self.assertTrue(issubclass(complex, Complex))
-        self.assertTrue(issubclass(complex, Number))
+        self.assertNotIsSubclass(complex, Integral)
+        self.assertNotIsSubclass(complex, Rational)
+        self.assertNotIsSubclass(complex, Real)
+        self.assertIsSubclass(complex, Complex)
+        self.assertIsSubclass(complex, Number)
 
         c1, c2 = complex(3, 2), complex(4,1)
         # XXX: This is not ideal, but see the comment in math_trunc().

Modules/_abc.c

@@ -578,6 +578,8 @@ _abc__abc_register_impl(PyObject *module, PyObject *self, PyObject *subclass)
     if (result < 0) {
         return NULL;
     }
+
+    /* Actual registration */
     _abc_data *impl = _get_impl(module, self);
     if (impl == NULL) {
         return NULL;
@@ -588,6 +590,49 @@ _abc__abc_register_impl(PyObject *module, PyObject *self, PyObject *subclass)
     }
     Py_DECREF(impl);
 
+    /* Recursively register the subclass in all ABC bases, to avoid recursive lookups.
+        >>> class Ancestor1(ABC): pass
+        >>> class Ancestor2(Ancestor1): pass
+        >>> class Other: pass
+        >>> Ancestor2.register(Other)  # same result for Ancestor1.register(Other)
+        >>> issubclass(Other, Ancestor2) is True
+        >>> issubclass(Other, Ancestor1) is True
+    */
+    PyObject *mro = PyObject_GetAttrString(self, "__mro__");
+    if (mro == NULL) {
+        return NULL;
+    }
+
+    if (!PyTuple_Check(mro)) {
+        PyErr_SetString(PyExc_TypeError, "__mro__ is not tuple");
+        goto error;
+    }
+
+    for (Py_ssize_t pos = 0; pos < PyTuple_GET_SIZE(mro); pos++) {
+        PyObject *base_class = PyTuple_GET_ITEM(mro, pos);  // borrowed
+        PyObject *base_class_data;
+
+        if (PyObject_GetOptionalAttr(base_class, &_Py_ID(_abc_impl),
+                                 &base_class_data) < 0) {
+            goto error;
+        }
+
+        if (PyErr_Occurred()) {
+            goto error;
+        }
+
+        if (base_class_data == NULL) {
+            // not ABC class
+            continue;
+        }
+
+        _abc_data *base_class_state = _abc_data_CAST(base_class_data);
+        if (_add_to_weak_set(base_class_state, &base_class_state->_abc_registry, subclass) < 0) {
+            Py_DECREF(base_class_data);
+            goto error;
+        }
+    }
+
     /* Invalidate negative cache */
     increment_invalidation_counter(get_abc_state(module));
 
@@ -602,6 +647,10 @@ _abc__abc_register_impl(PyObject *module, PyObject *self, PyObject *subclass)
         }
     }
     return Py_NewRef(subclass);
+
+error:
+    Py_XDECREF(mro);
+    return NULL;
 }
 
 
@@ -710,6 +759,7 @@ _abc__abc_subclasscheck_impl(PyObject *module, PyObject *self,
         PyErr_SetString(PyExc_TypeError, "issubclass() arg 1 must be a class");
         return NULL;
     }
+    PyTypeObject *cls = (PyTypeObject *)self;
 
     PyObject *ok, *subclasses = NULL, *result = NULL;
     _abcmodule_state *state = NULL;
@@ -800,32 +850,43 @@ _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) {
-            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;
-        }
+    /* 6. Check if it's a subclass of a subclass (recursive).
+        >>> class Ancestor: __subclasses__ = lambda: [Other]
+        >>> class Other: pass
+        >>> isinstance(Other, Ancestor) is True
+
+      Do not iterate over cls.__subclasses__() because it returns the entire class tree,
+      not just direct children, which leads to O(n^2) lookup.
+    */
+    PyObject *dict = _PyType_GetDict(cls); // borrowed
+    PyObject *subclasses_own_method = PyDict_GetItemString(dict, "__subclasses__"); // borrowed
+    if (subclasses_own_method) {
+      subclasses = PyObject_CallNoArgs(subclasses_own_method);
+      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) {
+              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;
+          }
+      }
     }
 
     /* No dice; update negative cache. */