Fully promote __init_subclass__ to a classmethod.

rchen152 · rchen152 · commit b9d11992d6ac · 2024-04-10T14:40:10.000-07:00
PiperOrigin-RevId: 623325031
diff --git a/pytype/abstract/_interpreter_function.py b/pytype/abstract/_interpreter_function.py
@@ -186,9 +186,6 @@ def __init__(self, name, def_opcode, code, f_locals, f_globals, defaults,
     self.last_frame = None  # for BuildClass
     self._store_call_records = False
     self.is_class_builder = False  # Will be set by BuildClass.
-    if name.endswith(".__init_subclass__"):
-      # __init_subclass__ is automatically promoted to a classmethod
-      self.is_classmethod = True
     # Whether to cache the return value irrespective of call args
     self.cache_return = False
 
diff --git a/pytype/tests/test_classes2.py b/pytype/tests/test_classes2.py
@@ -583,6 +583,16 @@ class X:
         x: int
     """)
 
+  def test_init_subclass_super(self):
+    self.Check("""
+      class A:
+        def __init_subclass__(cls):
+          pass
+      class B(A):
+        def __init_subclass__(cls):
+          super().__init_subclass__()
+    """)
+
 
 if __name__ == "__main__":
   test_base.main()
diff --git a/pytype/vm_utils.py b/pytype/vm_utils.py
@@ -526,6 +526,12 @@ def make_class(node, props, ctx):
             annotations_dict, ctx).to_variable(node)
         class_dict.members["__annotations__"] = annotations_member
         class_dict.pyval["__annotations__"] = annotations_member
+    if "__init_subclass__" in class_dict.members:
+      # __init_subclass__ is automatically promoted to a classmethod
+      underlying = class_dict.pyval["__init_subclass__"]
+      _, method = ctx.vm.load_special_builtin("classmethod").call(
+          node, func=None, args=function.Args(posargs=(underlying,)))
+      class_dict.pyval["__init_subclass__"] = method
     try:
       class_type = props.class_type or abstract.InterpreterClass
       assert issubclass(class_type, abstract.InterpreterClass)
