[mypyc] Introduce ClassBuilder and add support for attrs classes

mypyc/ir/class_ir.py

@@ -152,6 +152,14 @@ def __init__(self, name: str, module_name: str, is_trait: bool = False,
         # None if separate compilation prevents this from working
         self.children: Optional[List[ClassIR]] = []
 
+    def __repr__(self) -> str:
+        return (
+            "ClassIR("
+            "name={self.name}, module_name={self.module_name}, "
+            "is_trait={self.is_trait}, is_generated={self.is_generated}, "
+            "is_abstract={self.is_abstract}, is_ext_class={self.is_ext_class}"
+            ")".format(self=self))
+
     @property
     def fullname(self) -> str:
         return "{}.{}".format(self.module_name, self.name)

mypyc/irbuild/expression.py

@@ -38,7 +38,7 @@
 from mypyc.primitives.set_ops import set_add_op, set_update_op
 from mypyc.primitives.str_ops import str_slice_op
 from mypyc.primitives.int_ops import int_comparison_op_mapping
-from mypyc.irbuild.specialize import specializers
+from mypyc.irbuild.specialize import apply_function_specialization, apply_method_specialization
 from mypyc.irbuild.builder import IRBuilder
 from mypyc.irbuild.for_helpers import (
     translate_list_comprehension, translate_set_comprehension,
@@ -209,7 +209,8 @@ def transform_call_expr(builder: IRBuilder, expr: CallExpr) -> Value:
         callee = callee.analyzed.expr  # Unwrap type application
 
     if isinstance(callee, MemberExpr):
-        return translate_method_call(builder, expr, callee)
+        return apply_method_specialization(builder, expr, callee) or \
+               translate_method_call(builder, expr, callee)
     elif isinstance(callee, SuperExpr):
         return translate_super_method_call(builder, expr, callee)
     else:
@@ -219,7 +220,8 @@ def transform_call_expr(builder: IRBuilder, expr: CallExpr) -> Value:
 def translate_call(builder: IRBuilder, expr: CallExpr, callee: Expression) -> Value:
     # The common case of calls is refexprs
     if isinstance(callee, RefExpr):
-        return translate_refexpr_call(builder, expr, callee)
+        return apply_function_specialization(builder, expr, callee) or \
+               translate_refexpr_call(builder, expr, callee)
 
     function = builder.accept(callee)
     args = [builder.accept(arg) for arg in expr.args]
@@ -229,18 +231,6 @@ def translate_call(builder: IRBuilder, expr: CallExpr, callee: Expression) -> Va
 
 def translate_refexpr_call(builder: IRBuilder, expr: CallExpr, callee: RefExpr) -> Value:
     """Translate a non-method call."""
-
-    # TODO: Allow special cases to have default args or named args. Currently they don't since
-    # they check that everything in arg_kinds is ARG_POS.
-
-    # If there is a specializer for this function, try calling it.
-    # We would return the first successful one.
-    if callee.fullname and (callee.fullname, None) in specializers:
-        for specializer in specializers[callee.fullname, None]:
-            val = specializer(builder, expr, callee)
-            if val is not None:
-                return val
-
     # Gen the argument values
     arg_values = [builder.accept(arg) for arg in expr.args]
 
@@ -297,11 +287,9 @@ def translate_method_call(builder: IRBuilder, expr: CallExpr, callee: MemberExpr
 
         # If there is a specializer for this method name/type, try calling it.
         # We would return the first successful one.
-        if (callee.name, receiver_typ) in specializers:
-            for specializer in specializers[callee.name, receiver_typ]:
-                val = specializer(builder, expr, callee)
-                if val is not None:
-                    return val
+        val = apply_method_specialization(builder, expr, callee, receiver_typ)
+        if val is not None:
+            return val
 
         obj = builder.accept(callee.expr)
         args = [builder.accept(arg) for arg in expr.args]

mypyc/irbuild/function.py

@@ -117,16 +117,6 @@ def transform_decorator(builder: IRBuilder, dec: Decorator) -> None:
     builder.functions.append(func_ir)
 
 
-def transform_method(builder: IRBuilder,
-                     cdef: ClassDef,
-                     non_ext: Optional[NonExtClassInfo],
-                     fdef: FuncDef) -> None:
-    if non_ext:
-        handle_non_ext_method(builder, non_ext, cdef, fdef)
-    else:
-        handle_ext_method(builder, cdef, fdef)
-
-
 def transform_lambda_expr(builder: IRBuilder, expr: LambdaExpr) -> Value:
     typ = get_proper_type(builder.types[expr])
     assert isinstance(typ, CallableType)

mypyc/irbuild/specialize.py

@@ -57,6 +57,34 @@
 specializers: Dict[Tuple[str, Optional[RType]], List[Specializer]] = {}
 
 
+def _apply_specialization(builder: 'IRBuilder', expr: CallExpr, callee: RefExpr,
+                          name: Optional[str], typ: Optional[RType] = None) -> Optional[Value]:
+    # TODO: Allow special cases to have default args or named args. Currently they don't since
+    #       they check that everything in arg_kinds is ARG_POS.
+
+    # If there is a specializer for this function, try calling it.
+    # Return the first successful one.
+    if name and (name, typ) in specializers:
+        for specializer in specializers[name, typ]:
+            val = specializer(builder, expr, callee)
+            if val is not None:
+                return val
+    return None
+
+
+def apply_function_specialization(builder: 'IRBuilder', expr: CallExpr,
+                                  callee: RefExpr) -> Optional[Value]:
+    """Invoke the Specializer callback for a function if one has been registered"""
+    return _apply_specialization(builder, expr, callee, callee.fullname)
+
+
+def apply_method_specialization(builder: 'IRBuilder', expr: CallExpr, callee: MemberExpr,
+                                typ: Optional[RType] = None) -> Optional[Value]:
+    """Invoke the Specializer callback for a method if one has been registered"""
+    name = callee.fullname if typ is None else callee.name
+    return _apply_specialization(builder, expr, callee, name, typ)
+
+
 def specialize_function(
         name: str, typ: Optional[RType] = None) -> Callable[[Specializer], Specializer]:
     """Decorator to register a function as being a specializer.
@@ -267,10 +295,12 @@ def gen_inner_stmts() -> None:
 
 
 @specialize_function('dataclasses.field')
+@specialize_function('attr.ib')
+@specialize_function('attr.attrib')
 @specialize_function('attr.Factory')
 def translate_dataclasses_field_call(
         builder: IRBuilder, expr: CallExpr, callee: RefExpr) -> Optional[Value]:
-    """Special case for 'dataclasses.field' and 'attr.Factory'
+    """Special case for 'dataclasses.field', 'attr.attrib', and 'attr.Factory'
     function calls because the results of such calls are type-checked
     by mypy using the types of the arguments to their respective
     functions, resulting in attempted coercions by mypyc that throw a

mypyc/irbuild/util.py

@@ -4,11 +4,18 @@
 
 from mypy.nodes import (
     ClassDef, FuncDef, Decorator, OverloadedFuncDef, StrExpr, CallExpr, RefExpr, Expression,
-    IntExpr, FloatExpr, Var, TupleExpr, UnaryExpr, BytesExpr,
+    IntExpr, FloatExpr, Var, NameExpr, TupleExpr, UnaryExpr, BytesExpr,
     ArgKind, ARG_NAMED, ARG_NAMED_OPT, ARG_POS, ARG_OPT, GDEF,
 )
 
 
+DATACLASS_DECORATORS = {
+    'dataclasses.dataclass',
+    'attr.s',
+    'attr.attrs',
+}
+
+
 def is_trait_decorator(d: Expression) -> bool:
     return isinstance(d, RefExpr) and d.fullname == 'mypy_extensions.trait'
 
@@ -17,21 +24,40 @@ def is_trait(cdef: ClassDef) -> bool:
     return any(is_trait_decorator(d) for d in cdef.decorators) or cdef.info.is_protocol
 
 
-def is_dataclass_decorator(d: Expression) -> bool:
-    return (
-        (isinstance(d, RefExpr) and d.fullname == 'dataclasses.dataclass')
-        or (
-            isinstance(d, CallExpr)
+def dataclass_decorator_type(d: Expression) -> Optional[str]:
+    if isinstance(d, RefExpr) and d.fullname in DATACLASS_DECORATORS:
+        return d.fullname.split('.')[0]
+    elif (isinstance(d, CallExpr)
             and isinstance(d.callee, RefExpr)
-            and d.callee.fullname == 'dataclasses.dataclass'
-        )
-    )
+            and d.callee.fullname in DATACLASS_DECORATORS):
+        name = d.callee.fullname.split('.')[0]
+        if name == 'attr' and 'auto_attribs' in d.arg_names:
+            # Note: the mypy attrs plugin checks that the value of auto_attribs is
+            # not computed at runtime, so we don't need to perform that check here
+            auto = d.args[d.arg_names.index('auto_attribs')]
+            if isinstance(auto, NameExpr) and auto.name == 'True':
+                return 'attr-auto'
+        return name
+    else:
+        return None
+
+
+def is_dataclass_decorator(d: Expression) -> bool:
+    return dataclass_decorator_type(d) is not None
 
 
 def is_dataclass(cdef: ClassDef) -> bool:
     return any(is_dataclass_decorator(d) for d in cdef.decorators)
 
 
+def dataclass_type(cdef: ClassDef) -> Optional[str]:
+    for d in cdef.decorators:
+        typ = dataclass_decorator_type(d)
+        if typ is not None:
+            return typ
+    return None
+
+
 def get_mypyc_attr_literal(e: Expression) -> Any:
     """Convert an expression from a mypyc_attr decorator to a value.