[mypyc] Introduce ClassBuilder and add support for attrs classes (#11328)

Create a class structure for building different types of classes, extension vs non-extension, plus various types of dataclasses.

Specilize member expressions. Prior to this change dataclasses.field would only be specialized if called without the module, as field()

Add support for attrs classes and build proper __annotations__ for dataclasses.

Co-authored-by: Jingchen Ye <[email protected]>

Author: chadrik

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/classdef.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/expression.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/function.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.
@@ -329,10 +357,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/specialize.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.