cool fstrings

Author: pablogsal

Grammar/python.gram

@@ -966,9 +966,15 @@ fstring_middle[expr_ty]:
     | fstring_replacement_field
     | t=FSTRING_MIDDLE { _PyPegen_constant_from_token(p, t) }
 fstring_replacement_field[expr_ty]:
+    | '{' a=annotated_rhs debug_expr='='? conv=fstring_pretty_conversion ':' pretty_func=annotated_rhs rbrace='}' {
+        _PyPegen_pretty_formatted_value(p, a, debug_expr, conv, pretty_func, rbrace, EXTRA) }
+    | '{' a=annotated_rhs debug_expr='='? conv=fstring_pretty_conversion rbrace='}' {
+        _PyPegen_pretty_formatted_value(p, a, debug_expr, conv, NULL, rbrace, EXTRA) }
     | '{' a=annotated_rhs debug_expr='='? conversion=[fstring_conversion] format=[fstring_full_format_spec] rbrace='}' {
         _PyPegen_formatted_value(p, a, debug_expr, conversion, format, rbrace, EXTRA) }
     | invalid_fstring_replacement_field
+fstring_pretty_conversion[Token*]:
+    | conv_token="!" conv=NAME { _PyPegen_check_pretty_conversion(p, conv_token, conv) }
 fstring_conversion[ResultTokenWithMetadata*]:
     | conv_token="!" conv=NAME { _PyPegen_check_fstring_conversion(p, conv_token, conv) }
 fstring_full_format_spec[ResultTokenWithMetadata*]:

Include/ceval.h

@@ -125,13 +125,14 @@ PyAPI_FUNC(void) PyEval_ReleaseThread(PyThreadState *tstate);
                  }
 
 /* Masks and values used by FORMAT_VALUE opcode. */
-#define FVC_MASK      0x3
+#define FVC_MASK      0x7
 #define FVC_NONE      0x0
 #define FVC_STR       0x1
 #define FVC_REPR      0x2
 #define FVC_ASCII     0x3
-#define FVS_MASK      0x4
-#define FVS_HAVE_SPEC 0x4
+#define FVC_PRETTY    0x4
+#define FVS_MASK      0x8
+#define FVS_HAVE_SPEC 0x8
 
 #ifndef Py_LIMITED_API
 #  define Py_CPYTHON_CEVAL_H

Modules/_testinternalcapi/test_cases.c.h

@@ -212,9 +212,12 @@ _PyInterpolation_Build(PyObject *value, PyObject *str, int conversion, PyObject
             case FVC_STR:
                 interpolation->conversion = _Py_LATIN1_CHR('s');
                 break;
+            case FVC_PRETTY:
+                interpolation->conversion = _Py_LATIN1_CHR('p');
+                break;
             default:
                 PyErr_SetString(PyExc_SystemError,
-                    "Interpolation() argument 'conversion' must be one of 's', 'a' or 'r'");
+                    "Interpolation() argument 'conversion' must be one of 's', 'a', 'r', or 'p'");
                 Py_DECREF(interpolation);
                 return NULL;
         }

Objects/interpolationobject.c

@@ -1012,6 +1012,22 @@ _PyPegen_check_fstring_conversion(Parser *p, Token* conv_token, expr_ty conv)
     return result_token_with_metadata(p, conv, conv_token->metadata);
 }
 
+Token *
+_PyPegen_check_pretty_conversion(Parser *p, Token *conv_token, expr_ty conv)
+{
+    if (conv_token->lineno != conv->lineno || conv_token->end_col_offset != conv->col_offset) {
+        return NULL;
+    }
+    if (PyUnicode_GET_LENGTH(conv->v.Name.id) != 1) {
+        return NULL;
+    }
+    Py_UCS4 first = PyUnicode_READ_CHAR(conv->v.Name.id, 0);
+    if (first != 'p') {
+        return NULL;
+    }
+    return conv_token;
+}
+
 ResultTokenWithMetadata *
 _PyPegen_setup_full_format_spec(Parser *p, Token *colon, asdl_expr_seq *spec, int lineno, int col_offset,
                                 int end_lineno, int end_col_offset, PyArena *arena)
@@ -1633,6 +1649,38 @@ expr_ty _PyPegen_formatted_value(Parser *p, expr_ty expression, Token *debug, Re
     return _PyAST_JoinedStr(values, lineno, col_offset, debug_end_line, debug_end_offset, p->arena);
 }
 
+expr_ty _PyPegen_pretty_formatted_value(Parser *p, expr_ty expression, Token *debug,
+                                        Token *conv_token, expr_ty pretty_func,
+                                        Token *closing_brace, int lineno, int col_offset,
+                                        int end_lineno, int end_col_offset, PyArena *arena) {
+    int conversion_val = (int)'p';
+
+    expr_ty formatted_value = _PyAST_FormattedValue(
+        expression, conversion_val, pretty_func,
+        lineno, col_offset, end_lineno,
+        end_col_offset, arena
+    );
+
+    if (!debug) {
+        return formatted_value;
+    }
+
+    int debug_end_line = conv_token->lineno;
+    int debug_end_offset = conv_token->col_offset;
+    PyObject *debug_metadata = conv_token->metadata;
+
+    expr_ty debug_text = _PyAST_Constant(debug_metadata, NULL, lineno, col_offset + 1, debug_end_line,
+                                            debug_end_offset - 1, p->arena);
+    if (!debug_text) {
+        return NULL;
+    }
+
+    asdl_expr_seq *values = _Py_asdl_expr_seq_new(2, arena);
+    asdl_seq_SET(values, 0, debug_text);
+    asdl_seq_SET(values, 1, formatted_value);
+    return _PyAST_JoinedStr(values, lineno, col_offset, debug_end_line, debug_end_offset, p->arena);
+}
+
 static expr_ty
 _build_concatenated_bytes(Parser *p, asdl_expr_seq *strings, int lineno,
                         int col_offset, int end_lineno, int end_col_offset,

Parser/action_helpers.c

@@ -1099,6 +1099,7 @@ tok_get_normal_mode(struct tok_state *tok, tokenizer_mode* current_tok, struct t
         the_current_tok->last_expr_end = -1;
         the_current_tok->in_format_spec = 0;
         the_current_tok->in_debug = 0;
+        the_current_tok->in_pretty_conversion = 0;
 
         enum string_kind_t string_kind = FSTRING;
         switch (*tok->start) {
@@ -1268,7 +1269,20 @@ tok_get_normal_mode(struct tok_state *tok, tokenizer_mode* current_tok, struct t
             return MAKE_TOKEN(ERRORTOKEN);
         }
 
-        if (c == ':' && cursor == current_tok->curly_bracket_expr_start_depth) {
+        if (c == '!' && cursor == current_tok->curly_bracket_expr_start_depth) {
+            int peek1 = tok_nextc(tok);
+            if (peek1 == 'p') {
+                int peek2 = tok_nextc(tok);
+                if (peek2 == ':') {
+                    current_tok->in_pretty_conversion = 1;
+                }
+                tok_backup(tok, peek2);
+            }
+            tok_backup(tok, peek1);
+        }
+
+        if (c == ':' && cursor == current_tok->curly_bracket_expr_start_depth
+                && !current_tok->in_pretty_conversion) {
             current_tok->kind = TOK_FSTRING_MODE;
             current_tok->in_format_spec = 1;
             p_start = tok->start;
@@ -1368,6 +1382,7 @@ tok_get_normal_mode(struct tok_state *tok, tokenizer_mode* current_tok, struct t
                 current_tok->kind = TOK_FSTRING_MODE;
                 current_tok->in_format_spec = 0;
                 current_tok->in_debug = 0;
+                current_tok->in_pretty_conversion = 0;
             }
         }
         break;

Parser/lexer/lexer.c

@@ -66,6 +66,7 @@ typedef struct _tokenizer_mode {
     char* last_expr_buffer;
     int in_debug;
     int in_format_spec;
+    int in_pretty_conversion;
 
     enum string_kind_t string_kind;
 } tokenizer_mode;