py: Implement default and star args for lambdas.

3 files changed, 37 insertions, 29 deletions

diff --git a/py/compile.c b/py/compile.c
index 7ef25a7b90..10c3e2f483 100644
--- a/py/compile.c
+++ b/py/compile.c
@@ -572,8 +572,9 @@ STATIC void close_over_variables_etc(compiler_t *comp, scope_t *this_scope, int
     }
 }
 
-STATIC void compile_funcdef_param(compiler_t *comp, mp_parse_node_t pn) {
-    if (MP_PARSE_NODE_IS_STRUCT_KIND(pn, PN_typedargslist_star)) {
+STATIC void compile_funcdef_lambdef_param(compiler_t *comp, mp_parse_node_t pn) {
+    if (MP_PARSE_NODE_IS_STRUCT_KIND(pn, PN_typedargslist_star)
+        || MP_PARSE_NODE_IS_STRUCT_KIND(pn, PN_varargslist_star)) {
         comp->have_star = true;
         /* don't need to distinguish bare from named star
         mp_parse_node_struct_t *pns = (mp_parse_node_struct_t*)pn;
@@ -584,7 +585,8 @@ STATIC void compile_funcdef_param(compiler_t *comp, mp_parse_node_t pn) {
         }
         */
 
-    } else if (MP_PARSE_NODE_IS_STRUCT_KIND(pn, PN_typedargslist_dbl_star)) {
+    } else if (MP_PARSE_NODE_IS_STRUCT_KIND(pn, PN_typedargslist_dbl_star)
+        || MP_PARSE_NODE_IS_STRUCT_KIND(pn, PN_varargslist_dbl_star)) {
         // named double star
         // TODO do we need to do anything with this?
 
@@ -599,15 +601,21 @@ STATIC void compile_funcdef_param(compiler_t *comp, mp_parse_node_t pn) {
             pn_colon = MP_PARSE_NODE_NULL;
             pn_equal = MP_PARSE_NODE_NULL;
 
-        } else {
+        } else if (MP_PARSE_NODE_IS_STRUCT_KIND(pn, PN_typedargslist_name)) {
             // this parameter has a colon and/or equal specifier
 
-            assert(MP_PARSE_NODE_IS_STRUCT_KIND(pn, PN_typedargslist_name)); // should be
-
             mp_parse_node_struct_t *pns = (mp_parse_node_struct_t*)pn;
             pn_id = pns->nodes[0];
             pn_colon = pns->nodes[1];
             pn_equal = pns->nodes[2];
+
+        } else {
+            assert(MP_PARSE_NODE_IS_STRUCT_KIND(pn, PN_varargslist_name)); // should be
+            // this parameter has an equal specifier
+
+            mp_parse_node_struct_t *pns = (mp_parse_node_struct_t*)pn;
+            pn_id = pns->nodes[0];
+            pn_equal = pns->nodes[1];
         }
 
         if (MP_PARSE_NODE_IS_NULL(pn_equal)) {
@@ -653,24 +661,15 @@ STATIC void compile_funcdef_param(compiler_t *comp, mp_parse_node_t pn) {
     }
 }
 
-// leaves function object on stack
-// returns function name
-STATIC qstr compile_funcdef_helper(compiler_t *comp, mp_parse_node_struct_t *pns, uint emit_options) {
-    if (comp->pass == MP_PASS_SCOPE) {
-        // create a new scope for this function
-        scope_t *s = scope_new_and_link(comp, SCOPE_FUNCTION, (mp_parse_node_t)pns, emit_options);
-        // store the function scope so the compiling function can use it at each pass
-        pns->nodes[4] = (mp_parse_node_t)s;
-    }
-
+STATIC void compile_funcdef_lambdef(compiler_t *comp, scope_t *scope, mp_parse_node_t pn_params, pn_kind_t pn_list_kind) {
     // compile default parameters
     comp->have_star = false;
     comp->num_dict_params = 0;
     comp->num_default_params = 0;
-    apply_to_single_or_list(comp, pns->nodes[1], PN_typedargslist, compile_funcdef_param);
+    apply_to_single_or_list(comp, pn_params, pn_list_kind, compile_funcdef_lambdef_param);
 
     if (comp->compile_error != MP_OBJ_NULL) {
-        return MP_QSTR_NULL;
+        return;
     }
 
     // in Micro Python we put the default positional parameters into a tuple using the bytecode
@@ -680,11 +679,25 @@ STATIC qstr compile_funcdef_helper(compiler_t *comp, mp_parse_node_struct_t *pns
         EMIT(load_null); // sentinel indicating empty default keyword args
     }
 
+    // make the function
+    close_over_variables_etc(comp, scope, comp->num_default_params, comp->num_dict_params);
+}
+
+// leaves function object on stack
+// returns function name
+STATIC qstr compile_funcdef_helper(compiler_t *comp, mp_parse_node_struct_t *pns, uint emit_options) {
+    if (comp->pass == MP_PASS_SCOPE) {
+        // create a new scope for this function
+        scope_t *s = scope_new_and_link(comp, SCOPE_FUNCTION, (mp_parse_node_t)pns, emit_options);
+        // store the function scope so the compiling function can use it at each pass
+        pns->nodes[4] = (mp_parse_node_t)s;
+    }
+
     // get the scope for this function
     scope_t *fscope = (scope_t*)pns->nodes[4];
 
-    // make the function
-    close_over_variables_etc(comp, fscope, comp->num_default_params, comp->num_dict_params);
+    // compile the function definition
+    compile_funcdef_lambdef(comp, fscope, pns->nodes[1], PN_typedargslist);
 
     // return its name (the 'f' in "def f(...):")
     return fscope->simple_name;
@@ -1762,10 +1775,6 @@ STATIC void compile_test_if_expr(compiler_t *comp, mp_parse_node_struct_t *pns)
 }
 
 STATIC void compile_lambdef(compiler_t *comp, mp_parse_node_struct_t *pns) {
-    // TODO default params etc for lambda; possibly just use funcdef code
-    //mp_parse_node_t pn_params = pns->nodes[0];
-    //mp_parse_node_t pn_body = pns->nodes[1];
-
     if (comp->pass == MP_PASS_SCOPE) {
         // create a new scope for this lambda
         scope_t *s = scope_new_and_link(comp, SCOPE_LAMBDA, (mp_parse_node_t)pns, comp->scope_cur->emit_options);
@@ -1776,8 +1785,8 @@ STATIC void compile_lambdef(compiler_t *comp, mp_parse_node_struct_t *pns) {
     // get the scope for this lambda
     scope_t *this_scope = (scope_t*)pns->nodes[2];
 
-    // make the lambda
-    close_over_variables_etc(comp, this_scope, 0, 0);
+    // compile the lambda definition
+    compile_funcdef_lambdef(comp, this_scope, pns->nodes[0], PN_varargslist);
 }
 
 STATIC void compile_or_and_test(compiler_t *comp, mp_parse_node_struct_t *pns, bool cond) {
diff --git a/py/grammar.h b/py/grammar.h
index 03b15992d5..b507132d05 100644
--- a/py/grammar.h
+++ b/py/grammar.h
@@ -72,10 +72,10 @@ DEF_RULE(tfpdef, nc, and(2), tok(NAME), opt_rule(typedargslist_colon))
 // TODO varargslist lets through more than is allowed
 DEF_RULE(varargslist, nc, list_with_end, rule(varargslist_item), tok(DEL_COMMA))
 DEF_RULE(varargslist_item, nc, or(3), rule(varargslist_name), rule(varargslist_star), rule(varargslist_dbl_star))
-DEF_RULE(varargslist_name, nc, and(2), tok(NAME), opt_rule(varargslist_equal))
+DEF_RULE(varargslist_name, nc, ident | and(2), tok(NAME), opt_rule(varargslist_equal))
 DEF_RULE(varargslist_star, nc, and(2), tok(OP_STAR), opt_rule(vfpdef))
 DEF_RULE(varargslist_dbl_star, nc, and(2), tok(OP_DBL_STAR), tok(NAME))
-DEF_RULE(varargslist_equal, nc, and(2), tok(DEL_EQUAL), rule(test))
+DEF_RULE(varargslist_equal, nc, ident | and(2), tok(DEL_EQUAL), rule(test))
 DEF_RULE(vfpdef, nc, ident | and(1), tok(NAME))
 
 // stmt: compound_stmt | simple_stmt
diff --git a/py/parse.c b/py/parse.c
index 54c58fbcca..c1947f76bb 100644
--- a/py/parse.c
+++ b/py/parse.c
@@ -853,7 +853,6 @@ mp_parse_tree_t mp_parse(mp_lexer_t *lex, mp_parse_input_kind_t input_kind) {
                 // rule should not be emitted if it has only 1 argument
                 // NOTE: can't set this flag for atom_paren because we need it
                 // to distinguish, for example, [a,b] from [(a,b)]
-                // TODO possibly set for: varargslist_name, varargslist_equal
                 if (rule->act & RULE_ACT_ALLOW_IDENT) {
                     emit_rule = false;
                 }