[mypyc] Respect declared metaclasses for non-extension classes (python#8119)

msullivan · web-flow · commit 7e65f1e4ed1c · 2019-12-10T11:26:20.000-08:00
diff --git a/mypy/semanal.py b/mypy/semanal.py
@@ -1578,7 +1578,7 @@ def update_metaclass(self, defn: ClassDef) -> None:
         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)
+                base_expr.accept(self)
                 if (base_expr.callee.fullname in {'six.with_metaclass',
                                                   'future.utils.with_metaclass',
                                                   'past.utils.with_metaclass'}
diff --git a/mypyc/genops.py b/mypyc/genops.py
@@ -1375,10 +1375,7 @@ def load_non_ext_class(self, ir: ClassIR, non_ext: NonExtClassInfo, line: int) -
 
         self.finish_non_ext_dict(non_ext, line)
 
-        metaclass = self.primitive_op(type_object_op, [], line)
-        metaclass = self.primitive_op(py_calc_meta_op, [metaclass, non_ext.bases], line)
-
-        class_type_obj = self.py_call(metaclass,
+        class_type_obj = self.py_call(non_ext.metaclass,
                                       [cls_name, non_ext.bases, non_ext.dict],
                                       line)
         return class_type_obj
@@ -1453,16 +1450,22 @@ def add_non_ext_class_attr(self, non_ext: NonExtClassInfo, lvalue: NameExpr,
             ):
                 attr_to_cache.append(lvalue)
 
-    def setup_non_ext_dict(self, cdef: ClassDef, bases: Value) -> Value:
+    def find_non_ext_metaclass(self, cdef: ClassDef, bases: Value) -> Value:
+        """Find the metaclass of a class from its defs and bases. """
+        if cdef.metaclass:
+            declared_metaclass = self.accept(cdef.metaclass)
+        else:
+            declared_metaclass = self.primitive_op(type_object_op, [], cdef.line)
+
+        return self.primitive_op(py_calc_meta_op, [declared_metaclass, bases], cdef.line)
+
+    def setup_non_ext_dict(self, cdef: ClassDef, metaclass: Value, bases: Value) -> Value:
         """
         Initialize the class dictionary for a non-extension class. This class dictionary
         is passed to the metaclass constructor.
         """
 
         # Check if the metaclass defines a __prepare__ method, and if so, call it.
-        metaclass = self.primitive_op(type_object_op, [], cdef.line)
-        metaclass = self.primitive_op(py_calc_meta_op, [metaclass, bases],
-                                      cdef.line)
         has_prepare = self.primitive_op(py_hasattr_op,
                                         [metaclass,
                                         self.load_static_unicode('__prepare__')], cdef.line)
@@ -1506,6 +1509,7 @@ def dataclass_non_ext_info(self, cdef: ClassDef) -> Optional[NonExtClassInfo]:
                 self.primitive_op(new_dict_op, [], cdef.line),
                 self.add(TupleSet([], cdef.line)),
                 self.primitive_op(new_dict_op, [], cdef.line),
+                self.primitive_op(type_object_op, [], cdef.line),
             )
         else:
             return None
@@ -1560,12 +1564,13 @@ def visit_class_def(self, cdef: ClassDef) -> None:
             dataclass_non_ext = self.dataclass_non_ext_info(cdef)
         else:
             non_ext_bases = self.populate_non_ext_bases(cdef)
-            non_ext_dict = self.setup_non_ext_dict(cdef, non_ext_bases)
+            non_ext_metaclass = self.find_non_ext_metaclass(cdef, non_ext_bases)
+            non_ext_dict = self.setup_non_ext_dict(cdef, non_ext_metaclass, non_ext_bases)
             # We populate __annotations__ for non-extension classes
             # because dataclasses uses it to determine which attributes to compute on.
             # TODO: Maybe generate more precise types for annotations
             non_ext_anns = self.primitive_op(new_dict_op, [], cdef.line)
-            non_ext = NonExtClassInfo(non_ext_dict, non_ext_bases, non_ext_anns)
+            non_ext = NonExtClassInfo(non_ext_dict, non_ext_bases, non_ext_anns, non_ext_metaclass)
             dataclass_non_ext = None
             type_obj = None
 
diff --git a/mypyc/ops.py b/mypyc/ops.py
@@ -2022,17 +2022,15 @@ class NonExtClassInfo:
     """Information needed to construct a non-extension class.
 
 
-    Includes the class dictionary, a tuple of base classes, and
-    the class annotations dictionary.
+    Includes the class dictionary, a tuple of base classes,
+    the class annotations dictionary, and the metaclass.
     """
 
-    def __init__(self,
-            non_ext_dict: Value,
-            non_ext_bases: Value,
-            non_ext_anns: Value) -> None:
-        self.dict = non_ext_dict
-        self.bases = non_ext_bases
-        self.anns = non_ext_anns
+    def __init__(self, dict: Value, bases: Value, anns: Value, metaclass: Value) -> None:
+        self.dict = dict
+        self.bases = bases
+        self.anns = anns
+        self.metaclass = metaclass
 
 
 LiteralsMap = Dict[Tuple[Type[object], Union[int, float, str, bytes, complex]], str]
diff --git a/mypyc/test-data/run-classes.test b/mypyc/test-data/run-classes.test
@@ -1105,6 +1105,40 @@ try:
 except TypeError as e:
     assert(str(e) == "mypyc classes can't have a metaclass")
 
+[case testMetaclass]
+from meta import Meta
+import six
+
+class Nothing1(metaclass=Meta):
+    pass
+
+def ident(x): return x
+
+@ident
+class Test:
+    pass
+
+class Nothing2(six.with_metaclass(Meta, Test)):
+    pass
+
+@six.add_metaclass(Meta)
+class Nothing3:
+    pass
+
+[file meta.py]
+from typing import Any
+class Meta(type):
+    def __new__(mcs, name, bases, dct):
+        dct['X'] = 10
+        return super().__new__(mcs, name, bases, dct)
+
+
+[file driver.py]
+from native import Nothing1, Nothing2, Nothing3
+assert Nothing1.X == 10
+assert Nothing2.X == 10
+assert Nothing3.X == 10
+
 [case testPickling]
 from mypy_extensions import trait
 from typing import Any, TypeVar, Generic
