From b3cd41dd4ba9d44910b7efc6c4a1fa5f531858d4 Mon Sep 17 00:00:00 2001
From: Jared Hancock <jared.hancock@centeredsolutions.com>
Date: Mon, 30 May 2022 10:22:57 -0400
Subject: [PATCH] py/lexer: Allow conversion specifiers in f-strings (e.g. !r).

PEP-498 allows for conversion specifiers like !r and !s to convert the
expression declared in braces to be passed through repr() and str()
respectively.

This updates the logic that detects the end of the expression to also stop
when it sees "![rs]" that is either at the end of the f-string or before
the ":" indicating the start of the format specifier. The "![rs]" is now
retained in the format string, whereas previously it stayed on the end
of the expression leading to a syntax error.

Previously: `f"{x!y:z}"` --> `"{:z}".format(x!y)`
Now: `f"{x!y:z}"` --> `"{!y:z}".format(x)`

Note that "!a" is not supported by `str.format` as MicroPython has no
`ascii()`, but now this will raise the correct error.

Updated cpydiff and added tests.

Signed-off-by: Jim Mussared <jim.mussared@gmail.com>
---
 py/lexer.c                         | 24 ++++++++++++++++++------
 tests/basics/string_fstring.py     | 19 +++++++++++++++++++
 tests/cpydiff/core_fstring_repr.py | 18 ++++--------------
 3 files changed, 41 insertions(+), 20 deletions(-)

diff --git a/py/lexer.c b/py/lexer.c
index e7d0e81440..b22cc2eaea 100644
--- a/py/lexer.c
+++ b/py/lexer.c
@@ -361,13 +361,25 @@ STATIC void parse_string_literal(mp_lexer_t *lex, bool is_raw, bool is_fstring)
                     vstr_add_byte(&lex->fstring_args, '(');
                     // remember the start of this argument (if we need it for f'{a=}').
                     size_t i = lex->fstring_args.len;
-                    // extract characters inside the { until we reach the
-                    // format specifier or closing }.
-                    // (MicroPython limitation) note: this is completely unaware of
-                    // Python syntax and will not handle any expression containing '}' or ':'.
-                    // e.g. f'{"}"}' or f'{foo({})}'.
+                    // Extract characters inside the { until the bracket level
+                    // is zero and we reach the conversion specifier '!',
+                    // format specifier ':', or closing '}'. The conversion
+                    // and format specifiers are left unchanged in the format
+                    // string to be handled by str.format.
+                    // (MicroPython limitation) note: this is completely
+                    // unaware of Python syntax and will not handle any
+                    // expression containing '}' or ':'. e.g. f'{"}"}' or f'
+                    // {foo({})}'. However, detection of the '!' will
+                    // specifically ensure that it's followed by [rs] and
+                    // then either the format specifier or the closing
+                    // brace. This allows the use of e.g. != in expressions.
                     unsigned int nested_bracket_level = 0;
-                    while (!is_end(lex) && (nested_bracket_level != 0 || !is_char_or(lex, ':', '}'))) {
+                    while (!is_end(lex) && (nested_bracket_level != 0
+                                            || !(is_char_or(lex, ':', '}')
+                                                 || (is_char(lex, '!')
+                                                     && is_char_following_or(lex, 'r', 's')
+                                                     && is_char_following_following_or(lex, ':', '}'))))
+                           ) {
                         unichar c = CUR_CHAR(lex);
                         if (c == '[' || c == '{') {
                             nested_bracket_level += 1;
diff --git a/tests/basics/string_fstring.py b/tests/basics/string_fstring.py
index 8907a5c478..1a3960680e 100644
--- a/tests/basics/string_fstring.py
+++ b/tests/basics/string_fstring.py
@@ -61,3 +61,22 @@ except (ValueError, SyntaxError):
 print(f"a {1,} b")
 print(f"a {x,y,} b")
 print(f"a {x,1} b")
+
+# f-strings with conversion specifiers (only support !r and !s).
+a = "123"
+print(f"{a!r}")
+print(f"{a!s}")
+try:
+    eval('print(f"{a!x}")')
+except (ValueError, SyntaxError):
+    # CPython detects this at compile time, MicroPython fails with ValueError
+    # when the str.format is executed.
+    print("ValueError")
+
+# Mixing conversion specifiers with formatting.
+print(f"{a!r:8s}")
+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}")
diff --git a/tests/cpydiff/core_fstring_repr.py b/tests/cpydiff/core_fstring_repr.py
index df80abf795..d37fb48db7 100644
--- a/tests/cpydiff/core_fstring_repr.py
+++ b/tests/cpydiff/core_fstring_repr.py
@@ -1,18 +1,8 @@
 """
 categories: Core
-description: f-strings don't support the !r, !s, and !a conversions
-cause: MicroPython is optimised for code space.
-workaround: Use repr(), str(), and ascii() explicitly.
+description: f-strings don't support !a conversions
+cause: MicropPython does not implement ascii()
+workaround: None
 """
 
-
-class X:
-    def __repr__(self):
-        return "repr"
-
-    def __str__(self):
-        return "str"
-
-
-print(f"{X()!r}")
-print(f"{X()!s}")
+f"{'unicode text'!a}"