customize_class_mro plugin hook (TODO: tests)

mypy/plugin.py

@@ -173,6 +173,9 @@ def get_base_class_hook(self, fullname: str
                             ) -> Optional[Callable[[ClassDefContext], None]]:
         return None
 
+    def customize_class_mro(self, ctx: ClassDefContext) -> None:
+        pass
+
 
 T = TypeVar('T')
 
@@ -229,6 +232,10 @@ def get_base_class_hook(self, fullname: str
                             ) -> Optional[Callable[[ClassDefContext], None]]:
         return self._find_hook(lambda plugin: plugin.get_base_class_hook(fullname))
 
+    def customize_class_mro(self, ctx: ClassDefContext) -> None:
+        for plugin in self._plugins:
+            plugin.customize_class_mro(ctx)
+
     def _find_hook(self, lookup: Callable[[Plugin], T]) -> Optional[T]:
         for plugin in self._plugins:
             hook = lookup(plugin)

mypy/semanal.py

@@ -1119,14 +1119,23 @@ def analyze_base_classes(self, defn: ClassDef) -> None:
             # Give it an MRO consisting of just the class itself and object.
             defn.info.mro = [defn.info, self.object_type().type]
             return
-        calculate_class_mro(defn, self.fail_blocker)
+        self.calculate_class_mro(defn)
         # If there are cyclic imports, we may be missing 'object' in
         # the MRO. Fix MRO if needed.
         if info.mro and info.mro[-1].fullname() != 'builtins.object':
             info.mro.append(self.object_type().type)
         if defn.info.is_enum and defn.type_vars:
             self.fail("Enum class cannot be generic", defn)
 
+    def calculate_class_mro(self, defn: ClassDef) -> None:
+        try:
+            defn.info.calculate_mro()
+        except MroError:
+            self.fail_blocker("Cannot determine consistent method resolution order "
+                 '(MRO) for "%s"' % defn.name, defn)
+            defn.info.mro = []
+        self.plugin.customize_class_mro(ClassDefContext(defn, Expression(), self))
+
     def update_metaclass(self, defn: ClassDef) -> None:
         """Lookup for special metaclass declarations, and update defn fields accordingly.
 
@@ -3928,15 +3937,6 @@ def refers_to_class_or_function(node: Expression) -> bool:
             isinstance(node.node, (TypeInfo, FuncDef, OverloadedFuncDef)))
 
 
-def calculate_class_mro(defn: ClassDef, fail: Callable[[str, Context], None]) -> None:
-    try:
-        defn.info.calculate_mro()
-    except MroError:
-        fail("Cannot determine consistent method resolution order "
-             '(MRO) for "%s"' % defn.name, defn)
-        defn.info.mro = []
-
-
 def find_duplicate(list: List[T]) -> Optional[T]:
     """If the list has duplicates, return one of the duplicates.
 

mypy/semanal_pass3.py

@@ -28,10 +28,11 @@
 from mypy.traverser import TraverserVisitor
 from mypy.typeanal import TypeAnalyserPass3, collect_any_types
 from mypy.typevars import has_no_typevars
-from mypy.semanal_shared import PRIORITY_FORWARD_REF, PRIORITY_TYPEVAR_VALUES
+from mypy.semanal_shared import (
+    PRIORITY_FORWARD_REF, PRIORITY_TYPEVAR_VALUES, SemanticAnalyzerPass2
+)
 from mypy.subtypes import is_subtype
 from mypy.sametypes import is_same_type
-import mypy.semanal
 
 
 class SemanticAnalyzerPass3(TraverserVisitor):
@@ -42,7 +43,7 @@ class SemanticAnalyzerPass3(TraverserVisitor):
     """
 
     def __init__(self, modules: Dict[str, MypyFile], errors: Errors,
-                 sem: 'mypy.semanal.SemanticAnalyzerPass2') -> None:
+                 sem: SemanticAnalyzerPass2) -> None:
         self.modules = modules
         self.errors = errors
         self.sem = sem
@@ -123,7 +124,7 @@ def visit_class_def(self, tdef: ClassDef) -> None:
         # import loop. (Only do so if we succeeded the first time.)
         if tdef.info.mro:
             tdef.info.mro = []  # Force recomputation
-            mypy.semanal.calculate_class_mro(tdef, self.fail_blocker)
+            self.sem.calculate_class_mro(tdef)
             if tdef.info.is_protocol:
                 add_protocol_members(tdef.info)
         if tdef.analyzed is not None:
@@ -215,7 +216,7 @@ def visit_assignment_stmt(self, s: AssignmentStmt) -> None:
                     self.analyze_info(analyzed.info)
                 if analyzed.info and analyzed.info.mro:
                     analyzed.info.mro = []  # Force recomputation
-                    mypy.semanal.calculate_class_mro(analyzed.info.defn, self.fail_blocker)
+                    self.sem.calculate_class_mro(analyzed.info.defn)
             if isinstance(analyzed, TypeVarExpr):
                 types = []
                 if analyzed.upper_bound: