py/parse: Generate const int variables for annotated assignments.

agatti · agatti · commit 9f923a582e97 · 2025-07-09T01:47:33.000+02:00
This commit extends the parser's behaviour to also recognise constant integer assignments even when the assignment itself has an "int" type annotation. For assignments like "var = const(1)" the compiler generates an optimised constant load bytecode instruction, whilst before this change assignments like "var: int = const(1)" would generate a sequence of bytecode instructions to create a new integer object, create a variable, and perform the assignment. This fixes #15608. Signed-off-by: Alessandro Gatti <a.gatti@frob.it>
diff --git a/py/parse.c b/py/parse.c
@@ -705,6 +705,15 @@ static bool fold_logical_constants(parser_t *parser, uint8_t rule_id, size_t *nu
     return false;
 }
 
+#if MICROPY_COMP_CONST
+static bool is_const_atom_expr_normal(mp_parse_node_t node) {
+    return MP_PARSE_NODE_IS_STRUCT_KIND(node, RULE_atom_expr_normal)
+           && MP_PARSE_NODE_IS_ID(((mp_parse_node_struct_t *)node)->nodes[0])
+           && MP_PARSE_NODE_LEAF_ARG(((mp_parse_node_struct_t *)node)->nodes[0]) == MP_QSTR_const
+           && MP_PARSE_NODE_IS_STRUCT_KIND(((mp_parse_node_struct_t *)node)->nodes[1], RULE_trailer_paren);
+}
+#endif
+
 static bool fold_constants(parser_t *parser, uint8_t rule_id, size_t num_args) {
     // this code does folding of arbitrary integer expressions, eg 1 + 2 * 3 + 4
     // it does not do partial folding, eg 1 + 2 + x -> 3 + x
@@ -797,21 +806,26 @@ static bool fold_constants(parser_t *parser, uint8_t rule_id, size_t num_args) {
         if (!MP_PARSE_NODE_IS_NULL(pn1)
             && !(MP_PARSE_NODE_IS_STRUCT_KIND(pn1, RULE_expr_stmt_augassign)
                  || MP_PARSE_NODE_IS_STRUCT_KIND(pn1, RULE_expr_stmt_assign_list))) {
-            // this node is of the form <x> = <y>
+            // this node is of the form <x> = <y> or <x> : int = <y>
             mp_parse_node_t pn0 = peek_result(parser, 1);
             if (MP_PARSE_NODE_IS_ID(pn0)
-                && MP_PARSE_NODE_IS_STRUCT_KIND(pn1, RULE_atom_expr_normal)
-                && MP_PARSE_NODE_IS_ID(((mp_parse_node_struct_t *)pn1)->nodes[0])
-                && MP_PARSE_NODE_LEAF_ARG(((mp_parse_node_struct_t *)pn1)->nodes[0]) == MP_QSTR_const
-                && MP_PARSE_NODE_IS_STRUCT_KIND(((mp_parse_node_struct_t *)pn1)->nodes[1], RULE_trailer_paren)
-                ) {
+                && ((MP_PARSE_NODE_IS_STRUCT_KIND(pn1, RULE_annassign) &&
+                     MP_PARSE_NODE_IS_ID(((mp_parse_node_struct_t *)pn1)->nodes[0]) &&
+                     MP_PARSE_NODE_LEAF_ARG(((mp_parse_node_struct_t *)pn1)->nodes[0]) == MP_QSTR_int &&
+                     is_const_atom_expr_normal(((mp_parse_node_struct_t *)pn1)->nodes[1]))
+                    || is_const_atom_expr_normal(pn1))) {
                 // code to assign dynamic constants: id = const(value)
 
                 // get the id
                 qstr id = MP_PARSE_NODE_LEAF_ARG(pn0);
 
                 // get the value
-                mp_parse_node_t pn_value = ((mp_parse_node_struct_t *)((mp_parse_node_struct_t *)pn1)->nodes[1])->nodes[0];
+                mp_parse_node_t pn_value;
+                if (MP_PARSE_NODE_IS_STRUCT_KIND(pn1, RULE_annassign)) {
+                    pn_value = ((mp_parse_node_struct_t *)(((mp_parse_node_struct_t *)(((mp_parse_node_struct_t *)pn1)->nodes[1]))->nodes[1]))->nodes[0];
+                } else {
+                    pn_value = ((mp_parse_node_struct_t *)((mp_parse_node_struct_t *)pn1)->nodes[1])->nodes[0];
+                }
                 if (!mp_parse_node_is_const(pn_value)) {
                     mp_obj_t exc = mp_obj_new_exception_msg(&mp_type_SyntaxError,
                         MP_ERROR_TEXT("not a constant"));
