py/lexer: Support nested f-strings.

dpgeorge · dpgeorge · commit b51b1028e196 · 2025-12-19T00:27:27.000+11:00
Signed-off-by: Damien George &lt;damien@micropython.org&gt;
diff --git a/py/lexer.c b/py/lexer.c
@@ -148,25 +148,25 @@ static void next_char(mp_lexer_t *lex) {
     lex->chr0 = lex->chr1;
     lex->chr1 = lex->chr2;
 
-    // and add the next byte from either the fstring args or the reader
+    // and add the next byte from either inject_chrs or the reader
     #if MICROPY_PY_FSTRINGS
-    if (lex->fstring_args_idx) {
-        // if there are saved chars, then we're currently injecting fstring args
-        if (lex->fstring_args_idx < lex->fstring_args.len) {
-            lex->chr2 = lex->fstring_args.buf[lex->fstring_args_idx++];
+    if (lex->inject_chrs_idx) {
+        // if there are saved chars, then we're currently injecting them
+        if (lex->inject_chrs_idx < lex->inject_chrs.len) {
+            lex->chr2 = lex->inject_chrs.buf[lex->inject_chrs_idx++];
         } else {
-            // no more fstring arg bytes
+            // no more characters to inject
             lex->chr2 = '\0';
         }
 
         if (lex->chr0 == '\0') {
-            // consumed all fstring data, restore saved input queue
+            // consumed all injected characters, restore saved input queue
             lex->chr0 = lex->chr0_saved;
             lex->chr1 = lex->chr1_saved;
             lex->chr2 = lex->chr2_saved;
-            // stop consuming fstring arg data
-            vstr_reset(&lex->fstring_args);
-            lex->fstring_args_idx = 0;
+            // stop consuming injected characters
+            vstr_reset(&lex->inject_chrs);
+            lex->inject_chrs_idx = 0;
         }
     } else
     #endif
@@ -333,8 +333,7 @@ static void parse_string_literal(mp_lexer_t *lex, bool is_raw, bool is_fstring)
     #if MICROPY_PY_FSTRINGS
     if (is_fstring) {
         // assume there's going to be interpolation, so prep the injection data
-        // fstring_args_idx==0 && len(fstring_args)>0 means we're extracting the args.
-        // only when fstring_args_idx>0 will we consume the arg data
+        // len(fstring_args)>0 means we're extracting the args.
         // lex->fstring_args is reset when finished, so at this point there are two cases:
         // - lex->fstring_args is empty: start of a new f-string
         // - lex->fstring_args is non-empty: concatenation of adjacent f-strings
@@ -560,19 +559,25 @@ static bool skip_whitespace(mp_lexer_t *lex, bool stop_at_newline) {
 
 void mp_lexer_to_next(mp_lexer_t *lex) {
     #if MICROPY_PY_FSTRINGS
-    if (lex->fstring_args.len && lex->fstring_args_idx == 0) {
+    if (lex->fstring_args.len) {
         // moving onto the next token means the literal string is complete.
         // switch into injecting the format args.
         vstr_add_byte(&lex->fstring_args, ')');
-        lex->chr0_saved = lex->chr0;
-        lex->chr1_saved = lex->chr1;
-        lex->chr2_saved = lex->chr2;
-        lex->chr0 = lex->fstring_args.buf[0];
-        lex->chr1 = lex->fstring_args.buf[1];
-        lex->chr2 = lex->fstring_args.buf[2];
-        // we've already extracted 3 chars, but setting this non-zero also
-        // means we'll start consuming the fstring data
-        lex->fstring_args_idx = 3;
+        if (lex->inject_chrs_idx == 0) {
+            // switch from stream to inject_chrs
+            lex->chr0_saved = lex->chr0;
+            lex->chr1_saved = lex->chr1;
+            lex->chr2_saved = lex->chr2;
+        } else {
+            // already consuming from inject_chrs, rewind cached chars to insert new ones
+            assert(lex->inject_chrs_idx >= 3);
+            lex->inject_chrs_idx -= 3;
+        }
+        vstr_ins_strn(&lex->inject_chrs, lex->inject_chrs_idx, lex->fstring_args.buf, lex->fstring_args.len);
+        vstr_reset(&lex->fstring_args);
+        lex->chr0 = lex->inject_chrs.buf[lex->inject_chrs_idx++];
+        lex->chr1 = lex->inject_chrs.buf[lex->inject_chrs_idx++];
+        lex->chr2 = lex->inject_chrs.buf[lex->inject_chrs_idx++];
     }
     #endif
 
@@ -857,8 +862,9 @@ mp_lexer_t *mp_lexer_new(qstr src_name, mp_reader_t reader) {
     lex->indent_level = m_new(uint16_t, lex->alloc_indent_level);
     vstr_init(&lex->vstr, 32);
     #if MICROPY_PY_FSTRINGS
+    vstr_init(&lex->inject_chrs, 0);
+    lex->inject_chrs_idx = 0;
     vstr_init(&lex->fstring_args, 0);
-    lex->fstring_args_idx = 0;
     #endif
 
     // store sentinel for first indentation level
@@ -915,6 +921,7 @@ void mp_lexer_free(mp_lexer_t *lex) {
         lex->reader.close(lex->reader.data);
         vstr_clear(&lex->vstr);
         #if MICROPY_PY_FSTRINGS
+        vstr_clear(&lex->inject_chrs);
         vstr_clear(&lex->fstring_args);
         #endif
         m_del(uint16_t, lex->indent_level, lex->alloc_indent_level);
diff --git a/py/lexer.h b/py/lexer.h
@@ -182,8 +182,9 @@ typedef struct _mp_lexer_t {
     mp_token_kind_t tok_kind;   // token kind
     vstr_t vstr;                // token data
     #if MICROPY_PY_FSTRINGS
+    vstr_t inject_chrs;         // characters currently being injected into the stream
+    size_t inject_chrs_idx;     // current index into inject_chrs
     vstr_t fstring_args;        // extracted arguments to pass to .format()
-    size_t fstring_args_idx;    // how many bytes of fstring_args have been read
     #endif
 } mp_lexer_t;
 
diff --git a/tests/basics/string_fstring_nested.py b/tests/basics/string_fstring_nested.py
@@ -0,0 +1,9 @@
+# Test nesting of f-strings within f-strings.
+
+x = 1
+
+# 2-level nesting, with padding.
+print(f"a{f'b{x:2}c':>5}d")
+
+# 4-level nesting using the different styles of quotes.
+print(f"""a{f'''b{f"c{f'd{x}e'}f"}g'''}h""")
diff --git a/tests/basics/string_fstring_nested_py312.py b/tests/basics/string_fstring_nested_py312.py
@@ -0,0 +1,7 @@
+# Test nesting of f-strings within f-strings.
+# These test rely on Python 3.12+ to use the same quote style for nesting.
+
+x = 1
+
+# 8-level nesting using the same quote style.
+print(f"a{f"b{f"c{f"d{f"e{f"f{f"g{f"h{x}i"}j"}k"}l"}m"}n"}o"}p")
diff --git a/tests/basics/string_fstring_nested_py312.py.exp b/tests/basics/string_fstring_nested_py312.py.exp
@@ -0,0 +1 @@
+abcdefgh1ijklmnop
