py/lexer: Support concatenation of adjacent f-strings.

dpgeorge · dpgeorge · commit a066f2308f7b · 2024-06-06T14:58:46.000+10:00
This is quite a simple and small change to support concatenation of
adjacent f-strings, and improve compatibility with CPython.

Signed-off-by: Damien George &lt;damien@micropython.org&gt;
diff --git a/py/lexer.c b/py/lexer.c
@@ -336,8 +336,12 @@ static void parse_string_literal(mp_lexer_t *lex, bool is_raw, bool 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
-        // note: lex->fstring_args will be empty already (it's reset when finished)
-        vstr_add_str(&lex->fstring_args, ".format(");
+        // 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
+        if (vstr_len(&lex->fstring_args) == 0) {
+            vstr_add_str(&lex->fstring_args, ".format(");
+        }
     }
     #endif
 
diff --git a/tests/basics/string_fstring.py b/tests/basics/string_fstring.py
@@ -65,3 +65,13 @@ def foo(a, b):
 # Still allow ! in expressions.
 print(f"{'1' if a != '456' else '0'!r:8s}")
 print(f"{'1' if a != '456' else '0'!s:8s}")
+
+# Concatenation of adjacent f-strings.
+print(f"" f"")
+print(f"a" f"b")
+print(f"{x}" f"{y}")
+print(
+    f"a{x}b---------------------------------"
+    f"cd---------------------------------"
+    f"e{y}f---------------------------------"
+)
diff --git a/tests/cpydiff/core_fstring_concat.py b/tests/cpydiff/core_fstring_concat.py
@@ -1,13 +1,12 @@
 """
 categories: Core
-description: f-strings don't support concatenation with adjacent literals if the adjacent literals contain braces or are f-strings
+description: f-strings don't support concatenation with adjacent literals if the adjacent literals contain braces
 cause: MicroPython is optimised for code space.
-workaround: Use the + operator between literal strings when either or both are f-strings
+workaround: Use the + operator between literal strings when they are not both f-strings
 """
 
 x, y = 1, 2
 print("aa" f"{x}")  # works
 print(f"{x}" "ab")  # works
 print("a{}a" f"{x}")  # fails
 print(f"{x}" "a{}b")  # fails
-print(f"{x}" f"{y}")  # fails
