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

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 <damien@micropython.org>

3 files changed, 18 insertions, 5 deletions

diff --git a/py/lexer.c b/py/lexer.c
index bff8e6376..2774759a1 100644
--- 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
index 42d093b37..3a8969272 100644
--- a/tests/basics/string_fstring.py
+++ b/tests/basics/string_fstring.py
@@ -65,3 +65,13 @@ print(f"{a!s:8s}")
 # 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
index c2bdb4e66..83bf18820 100644
--- a/tests/cpydiff/core_fstring_concat.py
+++ b/tests/cpydiff/core_fstring_concat.py
@@ -1,8 +1,8 @@
 """
 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
@@ -10,4 +10,3 @@ 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