py/parse: Recognise const int assignments with type hints.

agatti · agatti · commit ef3ef6f4f89a · 2025-07-21T01:06:04.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. Now declarations like "var: int = const(val)" can be treated as constants by the compiler and thus be folded in the rest of the program as it sees fit. Before this change, that line would generate a variable creation and its value assignment, without any folding being performed. 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"));
diff --git a/tests/micropython/const2_annotated.py b/tests/micropython/const2_annotated.py
@@ -0,0 +1,48 @@
+# check that consts are not replaced in anything except standalone identifiers
+
+from micropython import const
+
+X: int = const(1)
+Y: int = const(2)
+Z: int = const(3)
+
+# import that uses a constant
+import micropython as X
+
+print(globals()["X"])
+
+
+# function name that matches a constant
+def X():
+    print("function X", X)
+
+
+globals()["X"]()
+
+
+# arguments that match a constant
+def f(X, *Y, **Z):
+    pass
+
+
+f(1)
+
+
+# class name that matches a constant
+class X:
+    def f(self):
+        print("class X", X)
+
+
+globals()["X"]().f()
+
+
+# constant within a class
+class A:
+    C1: int = const(4)
+
+    def X(self):
+        print("method X", Y, C1, self.C1)
+
+
+A().X()
diff --git a/tests/micropython/const2_annotated.py.exp b/tests/micropython/const2_annotated.py.exp
@@ -0,0 +1,4 @@
+<module 'micropython'>
+function X 1
+class X 1
+method X 2 4 4
diff --git a/tests/micropython/const_annotated.py b/tests/micropython/const_annotated.py
@@ -0,0 +1,30 @@
+# test constant optimisation
+# This test will only work when MICROPY_COMP_CONST is enabled.
+
+from micropython import const
+
+X: int = const(123)
+Y: int = const(X + 456)
+
+print(X, Y + 1)
+
+
+def f():
+    print(X, Y + 1)
+
+
+f()
+
+_X: int = const(12)
+_Y: int = const(_X + 34)
+
+print(_X, _Y)
+
+
+class A:
+    Z: int = const(1)
+    _Z: int = const(2)
+    print(Z, _Z)
+
+
+print(hasattr(A, "Z"), hasattr(A, "_Z"))
diff --git a/tests/micropython/const_annotated.py.exp b/tests/micropython/const_annotated.py.exp
@@ -0,0 +1,5 @@
+123 580
+123 580
+12 46
+1 2
+True False

-Original file line number
+Diff line change
@@ @@ -0,0 +1,5 @@ @@
 +123 580
 +123 580
 +12 46
 +1 2
 +True False