micropython · dpgeorge · Jun 16, 2020 · Jun 16, 2020 · Jun 16, 2020 · Jun 16, 2020
diff --git a/ports/bare-arm/mpconfigport.h b/ports/bare-arm/mpconfigport.h
@@ -21,6 +21,7 @@
 #define MICROPY_ERROR_REPORTING     (MICROPY_ERROR_REPORTING_TERSE)
 #define MICROPY_BUILTIN_METHOD_CHECK_SELF_ARG (0)
 #define MICROPY_PY_ASYNC_AWAIT (0)
+#define MICROPY_PY_ASSIGN_EXPR      (0)
 #define MICROPY_PY_BUILTINS_BYTEARRAY (0)
 #define MICROPY_PY_BUILTINS_DICT_FROMKEYS (0)
 #define MICROPY_PY_BUILTINS_MEMORYVIEW (0)

diff --git a/ports/minimal/mpconfigport.h b/ports/minimal/mpconfigport.h
@@ -19,6 +19,7 @@
 #define MICROPY_ERROR_REPORTING     (MICROPY_ERROR_REPORTING_TERSE)
 #define MICROPY_BUILTIN_METHOD_CHECK_SELF_ARG (0)
 #define MICROPY_PY_ASYNC_AWAIT      (0)
+#define MICROPY_PY_ASSIGN_EXPR      (0)
 #define MICROPY_PY_BUILTINS_BYTEARRAY (0)
 #define MICROPY_PY_BUILTINS_DICT_FROMKEYS (0)
 #define MICROPY_PY_BUILTINS_ENUMERATE (0)

diff --git a/py/compile.c b/py/compile.c
@@ -3,7 +3,7 @@
  *
  * The MIT License (MIT)
  *
- * Copyright (c) 2013-2015 Damien P. George
+ * Copyright (c) 2013-2020 Damien P. George
  *
  * Permission is hereby granted, free of charge, to any person obtaining a copy
  * of this software and associated documentation files (the "Software"), to deal
@@ -2108,6 +2108,27 @@ STATIC void compile_lambdef(compiler_t *comp, mp_parse_node_struct_t *pns) {
     compile_funcdef_lambdef(comp, this_scope, pns->nodes[0], PN_varargslist);
 }
 
+#if MICROPY_PY_ASSIGN_EXPR
+STATIC void compile_namedexpr_helper(compiler_t *comp, mp_parse_node_t pn_name, mp_parse_node_t pn_expr) {
+    if (!MP_PARSE_NODE_IS_ID(pn_name)) {
+        compile_syntax_error(comp, (mp_parse_node_t)pn_name, MP_ERROR_TEXT("can't assign to expression"));
+    }
+    compile_node(comp, pn_expr);
+    EMIT(dup_top);
+    scope_t *old_scope = comp->scope_cur;
+    if (SCOPE_IS_COMP_LIKE(comp->scope_cur->kind)) {
+        // Use parent's scope for assigned value so it can "escape"
+        comp->scope_cur = comp->scope_cur->parent;
+    }
+    compile_store_id(comp, MP_PARSE_NODE_LEAF_ARG(pn_name));
+    comp->scope_cur = old_scope;
+}
+
+STATIC void compile_namedexpr(compiler_t *comp, mp_parse_node_struct_t *pns) {
+    compile_namedexpr_helper(comp, pns->nodes[0], pns->nodes[1]);
+}
+#endif
+
 STATIC void compile_or_and_test(compiler_t *comp, mp_parse_node_struct_t *pns) {
     bool cond = MP_PARSE_NODE_STRUCT_KIND(pns) == PN_or_test;
     uint l_end = comp_next_label(comp);
@@ -2353,6 +2374,12 @@ STATIC void compile_trailer_paren_helper(compiler_t *comp, mp_parse_node_t pn_ar
                 star_flags |= MP_EMIT_STAR_FLAG_DOUBLE;
                 dblstar_args_node = pns_arg;
             } else if (MP_PARSE_NODE_STRUCT_KIND(pns_arg) == PN_argument) {
+                #if MICROPY_PY_ASSIGN_EXPR
+                if (MP_PARSE_NODE_IS_STRUCT_KIND(pns_arg->nodes[1], PN_argument_4)) {
+                    compile_namedexpr_helper(comp, pns_arg->nodes[0], ((mp_parse_node_struct_t *)pns_arg->nodes[1])->nodes[0]);
+                    n_positional++;
+                } else
+                #endif
                 if (!MP_PARSE_NODE_IS_STRUCT_KIND(pns_arg->nodes[1], PN_comp_for)) {
                     if (!MP_PARSE_NODE_IS_ID(pns_arg->nodes[0])) {
                         compile_syntax_error(comp, (mp_parse_node_t)pns_arg, MP_ERROR_TEXT("LHS of keyword arg must be an id"));
@@ -3083,7 +3110,7 @@ STATIC void compile_scope(compiler_t *comp, scope_t *scope, pass_kind_t pass) {
             EMIT_ARG(load_const_tok, MP_TOKEN_KW_NONE);
         }
         EMIT(return_value);
-    } else if (scope->kind == SCOPE_LIST_COMP || scope->kind == SCOPE_DICT_COMP || scope->kind == SCOPE_SET_COMP || scope->kind == SCOPE_GEN_EXPR) {
+    } else if (SCOPE_IS_COMP_LIKE(scope->kind)) {
         // a bit of a hack at the moment
 
         assert(MP_PARSE_NODE_IS_STRUCT(scope->pn));

diff --git a/py/grammar.h b/py/grammar.h
@@ -3,7 +3,7 @@
  *
  * The MIT License (MIT)
  *
- * Copyright (c) 2013-2015 Damien P. George
+ * Copyright (c) 2013-2020 Damien P. George
  *
  * Permission is hereby granted, free of charge, to any person obtaining a copy
  * of this software and associated documentation files (the "Software"), to deal
@@ -184,10 +184,10 @@ DEF_RULE_NC(async_stmt_2, or(3), rule(funcdef), rule(with_stmt), rule(for_stmt))
 #else
 DEF_RULE_NC(compound_stmt, or(8), rule(if_stmt), rule(while_stmt), rule(for_stmt), rule(try_stmt), rule(with_stmt), rule(funcdef), rule(classdef), rule(decorated))
 #endif
-DEF_RULE(if_stmt, c(if_stmt), and(6), tok(KW_IF), rule(test), tok(DEL_COLON), rule(suite), opt_rule(if_stmt_elif_list), opt_rule(else_stmt))
+DEF_RULE(if_stmt, c(if_stmt), and(6), tok(KW_IF), rule(namedexpr_test), tok(DEL_COLON), rule(suite), opt_rule(if_stmt_elif_list), opt_rule(else_stmt))
 DEF_RULE_NC(if_stmt_elif_list, one_or_more, rule(if_stmt_elif))
-DEF_RULE_NC(if_stmt_elif, and(4), tok(KW_ELIF), rule(test), tok(DEL_COLON), rule(suite))
-DEF_RULE(while_stmt, c(while_stmt), and(5), tok(KW_WHILE), rule(test), tok(DEL_COLON), rule(suite), opt_rule(else_stmt))
+DEF_RULE_NC(if_stmt_elif, and(4), tok(KW_ELIF), rule(namedexpr_test), tok(DEL_COLON), rule(suite))
+DEF_RULE(while_stmt, c(while_stmt), and(5), tok(KW_WHILE), rule(namedexpr_test), tok(DEL_COLON), rule(suite), opt_rule(else_stmt))
 DEF_RULE(for_stmt, c(for_stmt), and(7), tok(KW_FOR), rule(exprlist), tok(KW_IN), rule(testlist), tok(DEL_COLON), rule(suite), opt_rule(else_stmt))
 DEF_RULE(try_stmt, c(try_stmt), and(4), tok(KW_TRY), tok(DEL_COLON), rule(suite), rule(try_stmt_2))
 DEF_RULE_NC(try_stmt_2, or(2), rule(try_stmt_except_and_more), rule(try_stmt_finally))
@@ -210,6 +210,12 @@ DEF_RULE(suite_block_stmts, c(generic_all_nodes), one_or_more, rule(stmt))
 // lambdef: 'lambda' [varargslist] ':' test
 // lambdef_nocond: 'lambda' [varargslist] ':' test_nocond
 
+#if MICROPY_PY_ASSIGN_EXPR
+DEF_RULE(namedexpr_test, c(namedexpr), and_ident(2), rule(test), opt_rule(namedexpr_test_2))
+DEF_RULE_NC(namedexpr_test_2, and_ident(2), tok(OP_ASSIGN), rule(test))
+#else
+DEF_RULE_NC(namedexpr_test, or(1), rule(test))
+#endif
 DEF_RULE_NC(test, or(2), rule(lambdef), rule(test_if_expr))
 DEF_RULE(test_if_expr, c(test_if_expr), and_ident(2), rule(or_test), opt_rule(test_if_else))
 DEF_RULE_NC(test_if_else, and(4), tok(KW_IF), rule(or_test), tok(KW_ELSE), rule(test))
@@ -276,7 +282,7 @@ DEF_RULE_NC(atom_2b, or(2), rule(yield_expr), rule(testlist_comp))
 DEF_RULE(atom_bracket, c(atom_bracket), and(3), tok(DEL_BRACKET_OPEN), opt_rule(testlist_comp), tok(DEL_BRACKET_CLOSE))
 DEF_RULE(atom_brace, c(atom_brace), and(3), tok(DEL_BRACE_OPEN), opt_rule(dictorsetmaker), tok(DEL_BRACE_CLOSE))
 DEF_RULE_NC(testlist_comp, and_ident(2), rule(testlist_comp_2), opt_rule(testlist_comp_3))
-DEF_RULE_NC(testlist_comp_2, or(2), rule(star_expr), rule(test))
+DEF_RULE_NC(testlist_comp_2, or(2), rule(star_expr), rule(namedexpr_test))
 DEF_RULE_NC(testlist_comp_3, or(2), rule(comp_for), rule(testlist_comp_3b))
 DEF_RULE_NC(testlist_comp_3b, and_ident(2), tok(DEL_COMMA), opt_rule(testlist_comp_3c))
 DEF_RULE_NC(testlist_comp_3c, list_with_end, rule(testlist_comp_2), tok(DEL_COMMA))
@@ -342,7 +348,12 @@ DEF_RULE_NC(arglist_dbl_star, and(2), tok(OP_DBL_STAR), rule(test))
 // comp_if: 'if' test_nocond [comp_iter]
 
 DEF_RULE_NC(argument, and_ident(2), rule(test), opt_rule(argument_2))
+#if MICROPY_PY_ASSIGN_EXPR
+DEF_RULE_NC(argument_2, or(3), rule(comp_for), rule(argument_3), rule(argument_4))
+DEF_RULE_NC(argument_4, and(2), tok(OP_ASSIGN), rule(test))
+#else
 DEF_RULE_NC(argument_2, or(2), rule(comp_for), rule(argument_3))
+#endif
 DEF_RULE_NC(argument_3, and_ident(2), tok(DEL_EQUAL), rule(test))
 DEF_RULE_NC(comp_iter, or(2), rule(comp_for), rule(comp_if))
 DEF_RULE_NC(comp_for, and_blank(5), tok(KW_FOR), rule(exprlist), tok(KW_IN), rule(or_test), opt_rule(comp_iter))

diff --git a/py/lexer.c b/py/lexer.c
@@ -174,7 +174,8 @@ STATIC void indent_pop(mp_lexer_t *lex) {
 // this means if the start of two ops are the same then they are equal til the last char
 
 STATIC const char *const tok_enc =
-    "()[]{},:;~" // singles
+    "()[]{},;~"   // singles
+    ":e="         // : :=
     "<e=c<e="     // < <= << <<=
     ">e=c>e="     // > >= >> >>=
     "*e=c*e="     // * *= ** **=
@@ -194,8 +195,9 @@ STATIC const uint8_t tok_enc_kind[] = {
     MP_TOKEN_DEL_PAREN_OPEN, MP_TOKEN_DEL_PAREN_CLOSE,
     MP_TOKEN_DEL_BRACKET_OPEN, MP_TOKEN_DEL_BRACKET_CLOSE,
     MP_TOKEN_DEL_BRACE_OPEN, MP_TOKEN_DEL_BRACE_CLOSE,
-    MP_TOKEN_DEL_COMMA, MP_TOKEN_DEL_COLON, MP_TOKEN_DEL_SEMICOLON, MP_TOKEN_OP_TILDE,
+    MP_TOKEN_DEL_COMMA, MP_TOKEN_DEL_SEMICOLON, MP_TOKEN_OP_TILDE,
 
+    MP_TOKEN_DEL_COLON, MP_TOKEN_OP_ASSIGN,
     MP_TOKEN_OP_LESS, MP_TOKEN_OP_LESS_EQUAL, MP_TOKEN_OP_DBL_LESS, MP_TOKEN_DEL_DBL_LESS_EQUAL,
     MP_TOKEN_OP_MORE, MP_TOKEN_OP_MORE_EQUAL, MP_TOKEN_OP_DBL_MORE, MP_TOKEN_DEL_DBL_MORE_EQUAL,
     MP_TOKEN_OP_STAR, MP_TOKEN_DEL_STAR_EQUAL, MP_TOKEN_OP_DBL_STAR, MP_TOKEN_DEL_DBL_STAR_EQUAL,

diff --git a/py/lexer.h b/py/lexer.h
@@ -96,6 +96,7 @@ typedef enum _mp_token_kind_t {
     MP_TOKEN_KW_WITH,
     MP_TOKEN_KW_YIELD,
 
+    MP_TOKEN_OP_ASSIGN,
     MP_TOKEN_OP_TILDE,
 
     // Order of these 6 matches corresponding mp_binary_op_t operator

diff --git a/py/mpconfig.h b/py/mpconfig.h
@@ -835,6 +835,11 @@ typedef double mp_float_t;
 #define MICROPY_PY_ASYNC_AWAIT (1)
 #endif
 
+// Support for assignment expressions with := (see PEP 572, Python 3.8+)
+#ifndef MICROPY_PY_ASSIGN_EXPR
+#define MICROPY_PY_ASSIGN_EXPR (1)
+#endif
+
 // Non-standard .pend_throw() method for generators, allowing for
 // Future-like behavior with respect to exception handling: an
 // exception set with .pend_throw() will activate on the next call

diff --git a/py/scope.h b/py/scope.h
@@ -55,6 +55,7 @@ typedef struct _id_info_t {
 } id_info_t;
 
 #define SCOPE_IS_FUNC_LIKE(s) ((s) >= SCOPE_LAMBDA)
+#define SCOPE_IS_COMP_LIKE(s) (SCOPE_LIST_COMP <= (s) && (s) <= SCOPE_GEN_EXPR)
 
 // scope is a "block" in Python parlance
 typedef enum {

diff --git a/tests/basics/assign_expr.py b/tests/basics/assign_expr.py
@@ -0,0 +1,29 @@
+(x := 4)
+print(x)
+
+if x := 2:
+    print(True)
+print(x)
+
+print(4, x := 5)
+print(x)
+
+x = 1
+print(x, x := 5, x)
+print(x)
+
+
+def foo():
+    print("any", any((hit := i) % 5 == 3 and (hit % 2) == 0 for i in range(10)))
+    return hit
+
+
+hit = 123
+print(foo())
+print(hit)  # shouldn't be changed by foo
+
+print("any", any((hit := i) % 5 == 3 and (hit % 2) == 0 for i in range(10)))
+print(hit)  # should be changed by above
+
+print([((m := k + 1), k * m) for k in range(4)])
+print(m)
diff --git a/tests/basics/assign_expr.py.exp b/tests/basics/assign_expr.py.exp
@@ -0,0 +1,14 @@
+4
+True
+2
+4 5
+5
+1 5 5
+5
+any True
+8
+123
+any True
+8
+[(1, 0), (2, 2), (3, 6), (4, 12)]
+4
diff --git a/tests/basics/assign_expr_syntaxerror.py b/tests/basics/assign_expr_syntaxerror.py
@@ -0,0 +1,16 @@
+# test SyntaxError with := operator
+
+def test(code):
+    try:
+        print(eval(code))
+    except SyntaxError:
+        print('SyntaxError')
+
+test("x := 1")
+test("((x, y) := 1)")
+
+# these are currently all allowed in MicroPython, but not in CPython
+test("([i := i + 1 for i in range(4)])")
+test("([i := -1 for i, j in [(1, 2)]])")
+test("([[(i := j) for i in range(2)] for j in range(2)])")
+test("([[(j := i) for i in range(2)] for j in range(2)])")
diff --git a/tests/basics/assign_expr_syntaxerror.py.exp b/tests/basics/assign_expr_syntaxerror.py.exp
@@ -0,0 +1,6 @@
+SyntaxError
+SyntaxError
+[1, 2, 3, 4]
+[-1]
+[[0, 0], [1, 1]]
+[[0, 1], [0, 1]]
diff --git a/tests/cmdline/cmd_parsetree.py.exp b/tests/cmdline/cmd_parsetree.py.exp
@@ -3,7 +3,7 @@
          tok(4)
 [   4]   rule(22) (n=4)
            id(i)
-[   4]     rule(44) (n=1)
+[   4]     rule(45) (n=1)
              NULL
 [   5]     rule(8) (n=0)
            NULL

diff --git a/tests/cpydiff/syntax_assign_expr.py b/tests/cpydiff/syntax_assign_expr.py
@@ -0,0 +1,7 @@
+"""
+categories: Syntax,Operators
+description: MicroPython allows using := to assign to the variable of a comprehension, CPython raises a SyntaxError.
+cause: MicroPython is optimised for code size and doesn't check this case.
+workaround: Do not rely on this behaviour if writing CPython compatible code.
+"""
+print([i := -1 for i in range(4)])