Add Class Decorator/Metaclass/Base Class plugin

mypy/plugin.py

@@ -4,7 +4,7 @@
 from abc import abstractmethod
 from typing import Callable, List, Tuple, Optional, NamedTuple, TypeVar
 
-from mypy.nodes import Expression, StrExpr, IntExpr, UnaryExpr, Context, DictExpr
+from mypy.nodes import Expression, StrExpr, IntExpr, UnaryExpr, Context, DictExpr, ClassDef
 from mypy.types import (
     Type, Instance, CallableType, TypedDictType, UnionType, NoneTyp, FunctionLike, TypeVarType,
     AnyType, TypeList, UnboundType, TypeOfAny
@@ -53,6 +53,23 @@ def named_generic_type(self, name: str, args: List[Type]) -> Instance:
         raise NotImplementedError
 
 
+class SemanticAnalyzerPluginInterface:
+    """Interface for accessing semantic analyzer functionality in plugins."""
+
+    @abstractmethod
+    def named_type(self, qualified_name: str, args: Optional[List[Type]] = None) -> Instance:
+        raise NotImplementedError
+
+    @abstractmethod
+    def parse_bool(self, expr: Expression) -> Optional[bool]:
+        raise NotImplementedError
+
+    @abstractmethod
+    def fail(self, msg: str, ctx: Context, serious: bool = False, *,
+             blocker: bool = False) -> None:
+        raise NotImplementedError
+
+
 # A context for a function hook that infers the return type of a function with
 # a special signature.
 #
@@ -98,6 +115,14 @@ def named_generic_type(self, name: str, args: List[Type]) -> Instance:
         ('context', Context),
         ('api', CheckerPluginInterface)])
 
+# A context for a class hook that modifies the class definition.
+ClassDefContext = NamedTuple(
+    'ClassDecoratorContext', [
+        ('cls', ClassDef),       # The class definition
+        ('reason', Expression),  # The expression being applied (decorator, metaclass, base class)
+        ('api', SemanticAnalyzerPluginInterface)
+    ])
+
 
 class Plugin:
     """Base class of all type checker plugins.
@@ -136,7 +161,17 @@ def get_attribute_hook(self, fullname: str
                            ) -> Optional[Callable[[AttributeContext], Type]]:
         return None
 
-    # TODO: metaclass / class decorator hook
+    def get_class_decorator_hook(self, fullname: str
+                                 ) -> Optional[Callable[[ClassDefContext], None]]:
+        return None
+
+    def get_class_metaclass_hook(self, fullname: str
+                                 ) -> Optional[Callable[[ClassDefContext], None]]:
+        return None
+
+    def get_class_base_hook(self, fullname: str
+                            ) -> Optional[Callable[[ClassDefContext], None]]:
+        return None
 
 
 T = TypeVar('T')
@@ -182,6 +217,18 @@ def get_attribute_hook(self, fullname: str
                            ) -> Optional[Callable[[AttributeContext], Type]]:
         return self._find_hook(lambda plugin: plugin.get_attribute_hook(fullname))
 
+    def get_class_decorator_hook(self, fullname: str
+                                 ) -> Optional[Callable[[ClassDefContext], None]]:
+        return self._find_hook(lambda plugin: plugin.get_class_decorator_hook(fullname))
+
+    def get_class_metaclass_hook(self, fullname: str
+                                 ) -> Optional[Callable[[ClassDefContext], None]]:
+        return self._find_hook(lambda plugin: plugin.get_class_metaclass_hook(fullname))
+
+    def get_class_base_hook(self, fullname: str
+                            ) -> Optional[Callable[[ClassDefContext], None]]:
+        return self._find_hook(lambda plugin: plugin.get_class_base_hook(fullname))
+
     def _find_hook(self, lookup: Callable[[Plugin], T]) -> Optional[T]:
         for plugin in self._plugins:
             hook = lookup(plugin)

mypy/semanal.py

@@ -81,7 +81,7 @@
 from mypy.sametypes import is_same_type
 from mypy.options import Options
 from mypy import experiments
-from mypy.plugin import Plugin
+from mypy.plugin import Plugin, ClassDefContext, SemanticAnalyzerPluginInterface
 from mypy import join
 from mypy.util import get_prefix
 
@@ -172,7 +172,7 @@
 }
 
 
-class SemanticAnalyzerPass2(NodeVisitor[None]):
+class SemanticAnalyzerPass2(NodeVisitor[None], SemanticAnalyzerPluginInterface):
     """Semantically analyze parsed mypy files.
 
     The analyzer binds names and does various consistency checks for a
@@ -719,9 +719,44 @@ def analyze_class_body(self, defn: ClassDef) -> Iterator[bool]:
                 yield True
                 self.calculate_abstract_status(defn.info)
                 self.setup_type_promotion(defn)
-
+                self.apply_class_plugin_hooks(defn)
                 self.leave_class()
 
+    def apply_class_plugin_hooks(self, defn: ClassDef) -> None:
+        """Apply a plugin hook that may infer a more precise definition for a class."""
+        def get_fullname(expr: Expression) -> Optional[str]:
+            # We support @foo.bar(...) @foo.bar and @bar
+            # TODO: Support IndexExpressions?
+            if isinstance(expr, CallExpr):
+                if isinstance(expr.callee, RefExpr):
+                    return expr.callee.fullname
+            elif isinstance(expr, MemberExpr):
+                return get_member_expr_fullname(expr)
+            elif isinstance(expr, NameExpr):
+                return expr.fullname
+            return None
+
+        for decorator in defn.decorators:
+            decorator_name = get_fullname(decorator)
+            if decorator_name:
+                hook = self.plugin.get_class_decorator_hook(decorator_name)
+                if hook:
+                    hook(ClassDefContext(defn, decorator, self))
+
+        if defn.metaclass:
+            metaclass_name = get_fullname(defn.metaclass)
+            if metaclass_name:
+                hook = self.plugin.get_class_metaclass_hook(metaclass_name)
+                if hook:
+                    hook(ClassDefContext(defn, defn.metaclass, self))
+
+        for base in defn.base_type_exprs:
+            base_name = get_fullname(base)
+            if base_name:
+                hook = self.plugin.get_class_base_hook(base_name)
+                if hook:
+                    hook(ClassDefContext(defn, base, self))
+
     def analyze_class_keywords(self, defn: ClassDef) -> None:
         for value in defn.keywords.values():
             value.accept(self)