python · ilevkivskyi · Aug 31, 2017 · Aug 18, 2017 · Aug 18, 2017 · Aug 18, 2017
diff --git a/mypy/semanal.py b/mypy/semanal.py
@@ -679,6 +679,7 @@ def visit_class_def(self, defn: ClassDef) -> None:
     def analyze_class_body(self, defn: ClassDef) -> Iterator[bool]:
         with self.tvar_scope_frame(self.tvar_scope.class_frame()):
             is_protocol = self.detect_protocol_base(defn)
+            self.update_metaclass(defn)
             self.clean_up_bases_and_infer_type_variables(defn)
             self.analyze_class_keywords(defn)
             if self.analyze_typeddict_classdef(defn):
@@ -832,12 +833,8 @@ def clean_up_bases_and_infer_type_variables(self, defn: ClassDef) -> None:
         Now we will remove Generic[T] from bases of Foo and infer that the
         type variable 'T' is a type argument of Foo.
 
-        We also process six.with_metaclass() here.
-
         Note that this is performed *before* semantic analysis.
         """
-        # First process six.with_metaclass if present and well-formed
-        defn.base_type_exprs, defn.metaclass = self.check_with_metaclass(defn)
         removed = []  # type: List[int]
         declared_tvars = []  # type: TypeVarList
         for i, base_expr in enumerate(defn.base_type_exprs):
@@ -1090,18 +1087,56 @@ def analyze_base_classes(self, defn: ClassDef) -> None:
         if defn.info.is_enum and defn.type_vars:
             self.fail("Enum class cannot be generic", defn)
 
-    def check_with_metaclass(self, defn: ClassDef) -> Tuple[List[Expression],
-                                                            Optional[Expression]]:
-        # Special-case six.with_metaclass(M, B1, B2, ...).
-        if defn.metaclass is None and len(defn.base_type_exprs) == 1:
+    def update_metaclass(self, defn: ClassDef) -> None:
+        """Lookup for special metaclass declarations, and update defn fields accordingly.
+
+        * __metaclass__ attribute in Python 2
+        * six.with_metaclass(M, B1, B2, ...)
+        * @six.add_metaclass(M)
+        """
+
+        # Look for "__metaclass__ = <metaclass>" in Python 2
+        python2_meta_expr = None  # type: Optional[Expression]
+        if self.options.python_version[0] == 2:
+            for body_node in defn.defs.body:
+                if isinstance(body_node, ClassDef) and body_node.name == "__metaclass__":
+                    self.fail("Metaclasses defined as inner classes are not supported", body_node)
+                    break
+                elif isinstance(body_node, AssignmentStmt) and len(body_node.lvalues) == 1:
+                    lvalue = body_node.lvalues[0]
+                    if isinstance(lvalue, NameExpr) and lvalue.name == "__metaclass__":
+                        python2_meta_expr = body_node.rvalue
+
+        # Look for six.with_metaclass(M, B1, B2, ...)
+        with_meta_expr = None  # type: Optional[Expression]
+        if len(defn.base_type_exprs) == 1:
             base_expr = defn.base_type_exprs[0]
             if isinstance(base_expr, CallExpr) and isinstance(base_expr.callee, RefExpr):
                 base_expr.callee.accept(self)
                 if (base_expr.callee.fullname == 'six.with_metaclass'
                         and len(base_expr.args) >= 1
                         and all(kind == ARG_POS for kind in base_expr.arg_kinds)):
-                    return (base_expr.args[1:], base_expr.args[0])
-        return (defn.base_type_exprs, defn.metaclass)
+                    with_meta_expr = base_expr.args[0]
+                    defn.base_type_exprs = base_expr.args[1:]
+
+        # Look for @six.add_metaclass(M)
+        add_meta_expr = None  # type: Optional[Expression]
+        for dec_expr in defn.decorators:
+            if isinstance(dec_expr, CallExpr) and isinstance(dec_expr.callee, RefExpr):
+                dec_expr.callee.accept(self)
+                if (dec_expr.callee.fullname == 'six.add_metaclass'
+                    and len(dec_expr.args) == 1
+                        and dec_expr.arg_kinds[0] == ARG_POS):
+                    add_meta_expr = dec_expr.args[0]
+                    break
+
+        metas = {defn.metaclass, python2_meta_expr, with_meta_expr, add_meta_expr} - {None}
+        if len(metas) == 0:
+            return
+        if len(metas) > 1:
+            self.fail("Multiple metaclass definitions", defn)
+            return
+        defn.metaclass = metas.pop()
 
     def expr_to_analyzed_type(self, expr: Expression) -> Type:
         if isinstance(expr, CallExpr):
@@ -1150,16 +1185,6 @@ def is_base_class(self, t: TypeInfo, s: TypeInfo) -> bool:
         return False
 
     def analyze_metaclass(self, defn: ClassDef) -> None:
-        if defn.metaclass is None and self.options.python_version[0] == 2:
-            # Look for "__metaclass__ = <metaclass>" in Python 2.
-            for body_node in defn.defs.body:
-                if isinstance(body_node, ClassDef) and body_node.name == "__metaclass__":
-                    self.fail("Metaclasses defined as inner classes are not supported", body_node)
-                    return
-                elif isinstance(body_node, AssignmentStmt) and len(body_node.lvalues) == 1:
-                    lvalue = body_node.lvalues[0]
-                    if isinstance(lvalue, NameExpr) and lvalue.name == "__metaclass__":
-                        defn.metaclass = body_node.rvalue
         if defn.metaclass:
             if isinstance(defn.metaclass, NameExpr):
                 metaclass_name = defn.metaclass.name
@@ -1193,7 +1218,7 @@ def analyze_metaclass(self, defn: ClassDef) -> None:
         if defn.info.metaclass_type is None:
             # Inconsistency may happen due to multiple baseclasses even in classes that
             # do not declare explicit metaclass, but it's harder to catch at this stage
-            if defn.metaclass:
+            if defn.metaclass is not None:
                 self.fail("Inconsistent metaclass structure for '%s'" % defn.name, defn)
 
     def object_type(self) -> Instance:

diff --git a/test-data/unit/check-classes.test b/test-data/unit/check-classes.test
@@ -3470,64 +3470,94 @@ reveal_type(y['b']) # E: Revealed type is '__main__.B'
 -- Special support for six
 -- -----------------------
 
-[case testSixWithMetaclass]
+[case testSixMetaclass]
 import six
 class M(type):
     x = 5
 class A(six.with_metaclass(M)): pass
+@six.add_metaclass(M)
+class B: pass
 reveal_type(type(A).x)  # E: Revealed type is 'builtins.int'
+reveal_type(type(B).x)  # E: Revealed type is 'builtins.int'
 
-[case testSixWithMetaclass_python2]
+[case testSixMetaclass_python2]
 import six
 class M(type):
     x = 5
 class A(six.with_metaclass(M)): pass
+@six.add_metaclass(M)
+class B: pass
 reveal_type(type(A).x)  # E: Revealed type is 'builtins.int'
+reveal_type(type(B).x)  # E: Revealed type is 'builtins.int'
 
-[case testFromSixWithMetaclass]
-from six import with_metaclass
+[case testFromSixMetaclass]
+from six import with_metaclass, add_metaclass
 class M(type):
     x = 5
 class A(with_metaclass(M)): pass
+@add_metaclass(M)
+class B: pass
 reveal_type(type(A).x)  # E: Revealed type is 'builtins.int'
+reveal_type(type(B).x)  # E: Revealed type is 'builtins.int'
 
-[case testSixWithMetaclassImportFrom]
+[case testSixMetaclassImportFrom]
 import six
 from metadefs import M
 class A(six.with_metaclass(M)): pass
+@six.add_metaclass(M)
+class B: pass
 reveal_type(type(A).x)  # E: Revealed type is 'builtins.int'
+reveal_type(type(B).x)  # E: Revealed type is 'builtins.int'
 [file metadefs.py]
 class M(type):
     x = 5
 
-[case testSixWithMetaclassImport]
+[case testSixMetaclassImport]
 import six
 import metadefs
 class A(six.with_metaclass(metadefs.M)): pass
+@six.add_metaclass(metadefs.M)
+class B: pass
 reveal_type(type(A).x)  # E: Revealed type is 'builtins.int'
+reveal_type(type(B).x)  # E: Revealed type is 'builtins.int'
 [file metadefs.py]
 class M(type):
     x = 5
 
-[case testSixWithMetaclassAndBase]
+[case testSixMetaclassAndBase]
+from typing import Iterable, Iterator
 import six
-class M(type):
+class M(type, Iterable[int]):
     x = 5
+    def __iter__(self) -> Iterator[int]: ...
 class A:
     def foo(self): pass
 class B:
     def bar(self): pass
 class C1(six.with_metaclass(M, A)): pass
+@six.add_metaclass(M)
+class D1(A): pass
 class C2(six.with_metaclass(M, A, B)): pass
+@six.add_metaclass(M)
+class D2(A, B): pass
 reveal_type(type(C1).x)  # E: Revealed type is 'builtins.int'
+reveal_type(type(D1).x)  # E: Revealed type is 'builtins.int'
 reveal_type(type(C2).x)  # E: Revealed type is 'builtins.int'
+reveal_type(type(D2).x)  # E: Revealed type is 'builtins.int'
 C1().foo()
+D1().foo()
 C1().bar()  # E: "C1" has no attribute "bar"
+D1().bar()  # E: "D1" has no attribute "bar"
+for x in C1: reveal_type(x)  # E: Revealed type is 'builtins.int*'
+for x in C2: reveal_type(x)  # E: Revealed type is 'builtins.int*'
 C2().foo()
+D2().foo()
 C2().bar()
+D2().bar()
 C2().baz()  # E: "C2" has no attribute "baz"
+D2().baz()  # E: "D2" has no attribute "baz"
 
-[case testSixWithMetaclassGenerics]
+[case testSixMetaclassGenerics]
 from typing import Generic, GenericMeta, TypeVar
 import six
 class DestroyableMeta(type):
@@ -3539,25 +3569,52 @@ class ArcMeta(GenericMeta, DestroyableMeta):
     pass
 class Arc(six.with_metaclass(ArcMeta, Generic[T_co], Destroyable)):
     pass
+@six.add_metaclass(ArcMeta)
+class Arc1(Generic[T_co], Destroyable):
+    pass
 class MyDestr(Destroyable):
     pass
 reveal_type(Arc[MyDestr]())  # E: Revealed type is '__main__.Arc[__main__.MyDestr*]'
+reveal_type(Arc1[MyDestr]())  # E: Revealed type is '__main__.Arc1[__main__.MyDestr*]'
 [builtins fixtures/bool.pyi]
 
-[case testSixWithMetaclassErrors]
+[case testSixMetaclassErrors]
 import six
 class M(type): pass
 class A(object): pass
 def f() -> type: return M
 class C1(six.with_metaclass(M), object): pass  # E: Invalid base class
 class C2(C1, six.with_metaclass(M)): pass  # E: Invalid base class
 class C3(six.with_metaclass(A)): pass  # E: Metaclasses not inheriting from 'type' are not supported
-class C4(six.with_metaclass(M), metaclass=M): pass  # E: Invalid base class
+@six.add_metaclass(A)  # E: Metaclasses not inheriting from 'type' are not supported
+class D3(A): pass
+class C4(six.with_metaclass(M), metaclass=M): pass  # E: Multiple metaclass definitions
+@six.add_metaclass(M)  # E: Multiple metaclass definitions
+class D4(metaclass=M): pass
 class C5(six.with_metaclass(f())): pass  # E: Dynamic metaclass not supported for 'C5'
+@six.add_metaclass(f())  # E: Dynamic metaclass not supported for 'D5'
+class D5: pass
+
+@six.add_metaclass(M)  # E: Multiple metaclass definitions
+class CD(six.with_metaclass(M)): pass
 
-[case testSixWithMetaclassErrors_python2-skip]
-# No error here yet
+class M1(type): pass
+class Q1(metaclass=M1): pass
+@six.add_metaclass(M)  # E: Inconsistent metaclass structure for 'CQA'
+class CQA(Q1): pass
+class CQW(six.with_metaclass(M, Q1)): pass  # E: Inconsistent metaclass structure for 'CQW'
+
+[case testSixMetaclassErrors_python2]
+# flags: --python-version 2.7
 import six
 class M(type): pass
-class C4(six.with_metaclass(M)):
+class C4(six.with_metaclass(M)):  # E: Multiple metaclass definitions
     __metaclass__ = M
+
+[case testSixMetaclassAny]
+import t  # type: ignore
+import six
+class E(metaclass=t.M): pass
+class F(six.with_metaclass(t.M)): pass
+@six.add_metaclass(t.M)
+class G: pass
diff --git a/test-data/unit/lib-stub/six.pyi b/test-data/unit/lib-stub/six.pyi
@@ -1,2 +1,3 @@
-from typing import Type
+from typing import Type, Callable
 def with_metaclass(mcls: Type[type], *args: type) -> type: pass
+def add_metaclass(mcls: Type[type]) -> Callable[[type], type]: pass