Implement real support for default arguments (mypyc/mypyc#680)

msullivan · web-flow · commit 56ed4947c5d3 · 2019-08-08T15:53:10.000-07:00
Compute the defaults at function creation time and store them in globals (for toplevel functions) or in the function object (for nested functions). Fixes mypyc/mypyc#336. No tests for lambdas yet because they are broken due to a mypy bug I've sent a fix up for (python#7306). (Also fix a silly lambda bug where defaults weren't supported I came across while testing.)
diff --git a/mypyc/genops.py b/mypyc/genops.py
@@ -1176,7 +1176,7 @@ def visit_method(self, cdef: ClassDef, fdef: FuncDef) -> None:
         else:
             self.handle_non_ext_method(cdef, fdef)
 
-    def is_approximately_constant(self, e: Expression) -> bool:
+    def is_constant(self, e: Expression) -> bool:
         """Check whether we allow an expression to appear as a default value.
 
         We don't currently properly support storing the evaluated
@@ -1189,7 +1189,7 @@ def is_approximately_constant(self, e: Expression) -> bool:
                 or (isinstance(e, UnaryExpr) and e.op == '-'
                     and isinstance(e.expr, (IntExpr, FloatExpr)))
                 or (isinstance(e, TupleExpr)
-                    and all(self.is_approximately_constant(e) for e in e.items))
+                    and all(self.is_constant(e) for e in e.items))
                 or (isinstance(e, RefExpr) and e.kind == GDEF
                     and (e.fullname in ('builtins.True', 'builtins.False', 'builtins.None')
                          or (isinstance(e.node, Var) and e.node.is_final))))
@@ -1228,7 +1228,7 @@ def generate_attr_defaults(self, cdef: ClassDef) -> None:
         for stmt in default_assignments:
             lvalue = stmt.lvalues[0]
             assert isinstance(lvalue, NameExpr)
-            if not stmt.is_final_def and not self.is_approximately_constant(stmt.rvalue):
+            if not stmt.is_final_def and not self.is_constant(stmt.rvalue):
                 self.warning('Unsupported default attribute value', stmt.rvalue.line)
 
             # If the attribute is initialized to None and type isn't optional,
@@ -1693,20 +1693,61 @@ def create_ne_from_eq(self, cdef: ClassDef) -> None:
             cls.methods['__ne__'] = f
             self.functions.append(f)
 
-    def gen_arg_default(self) -> None:
+    def calculate_arg_defaults(self,
+                               fn_info: FuncInfo,
+                               env: Environment,
+                               func_reg: Optional[Value]) -> None:
+        """Calculate default argument values and store them.
+
+        They are stored in statics for top level functions and in
+        the function objects for nested functions (while constants are
+        still stored computed on demand).
+        """
+        fitem = fn_info.fitem
+        for arg in fitem.arguments:
+            # Constant values don't get stored but just recomputed
+            if arg.initializer and not self.is_constant(arg.initializer):
+                value = self.coerce(self.accept(arg.initializer),
+                                    env.lookup(arg.variable).type, arg.line)
+                if not fn_info.is_nested:
+                    name = fitem.fullname() + '.' + arg.variable.name()
+                    self.add(InitStatic(value, name, 'final'))
+                else:
+                    assert func_reg is not None
+                    self.add(SetAttr(func_reg, arg.variable.name(), value, arg.line))
+
+    def gen_arg_defaults(self) -> None:
         """Generate blocks for arguments that have default values.
 
-        If the passed value is an error value, then assign the default value to the argument.
+        If the passed value is an error value, then assign the default
+        value to the argument.
         """
         fitem = self.fn_info.fitem
         for arg in fitem.arguments:
             if arg.initializer:
-                if not self.is_approximately_constant(arg.initializer):
-                    self.warning('Unsupported default argument value', arg.initializer.line)
                 target = self.environment.lookup(arg.variable)
+
+                def get_default() -> Value:
+                    assert arg.initializer is not None
+
+                    # If it is constant, don't bother storing it
+                    if self.is_constant(arg.initializer):
+                        return self.accept(arg.initializer)
+
+                    # Because gen_arg_defaults runs before calculate_arg_defaults, we
+                    # add the static/attribute to final_names/the class here.
+                    elif not self.fn_info.is_nested:
+                        name = fitem.fullname() + '.' + arg.variable.name()
+                        self.final_names.append((name, target.type))
+                        return self.add(LoadStatic(target.type, name, 'final'))
+                    else:
+                        name = arg.variable.name()
+                        self.fn_info.callable_class.ir.attributes[name] = target.type
+                        return self.add(
+                            GetAttr(self.fn_info.callable_class.self_reg, name, arg.line))
                 assert isinstance(target, AssignmentTargetRegister)
                 self.assign_if_null(target,
-                                    lambda: self.accept(cast(Expression, arg.initializer)),
+                                    get_default,
                                     arg.initializer.line)
 
     def gen_func_item(self,
@@ -1793,7 +1834,7 @@ def c() -> None:
             self.add_raise_exception_blocks_to_generator_class(fitem.line)
         else:
             self.load_env_registers()
-            self.gen_arg_default()
+            self.gen_arg_defaults()
 
         if self.fn_info.contains_nested and not self.fn_info.is_generator:
             self.finalize_env_class()
@@ -1850,6 +1891,8 @@ def c() -> None:
         else:
             func_ir, func_reg = self.gen_func_ir(blocks, sig, env, fn_info, cdef)
 
+        self.calculate_arg_defaults(fn_info, env, func_reg)
+
         return (func_ir, func_reg)
 
     def gen_func_ir(self,
@@ -3776,7 +3819,8 @@ def visit_lambda_expr(self, expr: LambdaExpr) -> Value:
         runtime_args = []
         for arg, arg_type in zip(expr.arguments, typ.arg_types):
             arg.variable.type = arg_type
-            runtime_args.append(RuntimeArg(arg.variable.name(), self.type_to_rtype(arg_type)))
+            runtime_args.append(
+                RuntimeArg(arg.variable.name(), self.type_to_rtype(arg_type), arg.kind))
         ret_type = self.type_to_rtype(typ.ret_type)
 
         fsig = FuncSignature(runtime_args, ret_type)
@@ -4857,7 +4901,7 @@ def instantiate_env_class(self) -> Value:
     def gen_generator_func(self) -> None:
         self.setup_generator_class()
         self.load_env_registers()
-        self.gen_arg_default()
+        self.gen_arg_defaults()
         self.finalize_env_class()
         self.add(Return(self.instantiate_generator_class()))
 
diff --git a/test-data/commandline.test b/test-data/commandline.test
@@ -148,9 +148,6 @@ class Concrete2:
 class Trait2(Concrete2):
     pass
 
-def warning(x: List[int] = []) -> None:  # W: Unsupported default argument value
-    pass
-
 class Nope(Trait1, Concrete2):  # E: Non-trait bases must appear first in parent list  # E: Non-trait MRO must be linear
     pass
 
diff --git a/test-data/run.test b/test-data/run.test
@@ -635,7 +635,7 @@ with assertRaises(NameError):
 lol
 
 [case testFunctionCallWithDefaultArgs]
-from typing import Tuple, List, Optional
+from typing import Tuple, List, Optional, Callable, Any
 def f(x: int, y: int = 3, s: str = "test", z: object = 5) -> Tuple[int, str]:
     def inner() -> int:
         return x + y
@@ -645,15 +645,36 @@ def g() -> None:
     assert f(s = "123", x = -2) == (1, "123")
 def h(a: Optional[object] = None, b: Optional[str] = None) -> Tuple[object, Optional[str]]:
     return (a, b)
+
+def same(x: object = object()) -> object:
+    return x
+
+a_lambda: Callable[..., Any] = lambda n=20: n
+
+def nested_funcs(n: int) -> List[Callable[..., Any]]:
+    ls: List[Callable[..., Any]] = []
+    for i in range(n):
+        def f(i: int = i) -> int:
+            return i
+        ls.append(f)
+    return ls
+
+
 [file driver.py]
-from native import f, g, h
+from native import f, g, h, same, nested_funcs, a_lambda
 g()
 assert f(2) == (5, "test")
 assert f(s = "123", x = -2) == (1, "123")
 assert h() == (None, None)
 assert h(10) == (10, None)
 assert h(b='a') == (None, 'a')
 assert h(10, 'a') == (10, 'a')
+assert same() == same()
+
+assert [f() for f in nested_funcs(10)] == list(range(10))
+
+assert a_lambda(10) == 10
+assert a_lambda() == 20
 
 [case testMethodCallWithDefaultArgs]
 from typing import Tuple, List
