py/lexer: Support nested [] and {} characters within f-string params.

Signed-off-by: Damien George <damien@micropython.org>

py/lexer.c

@@ -363,9 +363,16 @@ STATIC void parse_string_literal(mp_lexer_t *lex, bool is_raw, bool is_fstring)
                     // (MicroPython limitation) note: this is completely unaware of
                     // Python syntax and will not handle any expression containing '}' or ':'.
                     // e.g. f'{"}"}' or f'{foo({})}'.
-                    while (!is_end(lex) && !is_char_or(lex, ':', '}')) {
+                    unsigned int nested_bracket_level = 0;
+                    while (!is_end(lex) && (nested_bracket_level != 0 || !is_char_or(lex, ':', '}'))) {
+                        unichar c = CUR_CHAR(lex);
+                        if (c == '[' || c == '{') {
+                            nested_bracket_level += 1;
+                        } else if (c == ']' || c == '}') {
+                            nested_bracket_level -= 1;
+                        }
                         // like the default case at the end of this function, stay 8-bit clean
-                        vstr_add_byte(&lex->fstring_args, CUR_CHAR(lex));
+                        vstr_add_byte(&lex->fstring_args, c);
                         next_char(lex);
                     }
                     if (lex->fstring_args.buf[lex->fstring_args.len - 1] == '=') {

tests/basics/string_fstring.py

@@ -22,6 +22,13 @@ def foo(a, b):
     return f'{x}{y}{a}{b}'
 print(foo(7, 8))
 
+# ':' character within {...} that should not be interpreted as format specifiers.
+print(f"a{[0,1,2][0:2]}")
+print(f"a{[0,15,2][0:2][-1]:04x}")
+
+# Nested '{' and '}' characters.
+print(f"a{ {0,1,2}}")
+
 # PEP-0498 specifies that '\\' and '#' must be disallowed explicitly, whereas
 # MicroPython relies on the syntax error as a result of the substitution.
 

tests/cpydiff/core_fstring_parser.py

 """
 categories: Core
-description: f-strings cannot support expressions that require parsing to resolve nested braces
+description: f-strings cannot support expressions that require parsing to resolve unbalanced nested braces and brackets
 cause: MicroPython is optimised for code space.
-workaround: Only use simple expressions inside f-strings
+workaround: Always use balanced braces and brackets in expressions inside f-strings
 """
 
-f'{"hello {} world"}'
-f"{repr({})}"
+print(f'{"hello { world"}')
+print(f'{"hello ] world"}')