gh-132661: Implement PEP 750 (#132662)

Co-authored-by: Lysandros Nikolaou <lisandrosnik@gmail.com>
Co-authored-by: Bénédikt Tran <10796600+picnixz@users.noreply.github.com>
Co-authored-by: Adam Turner <9087854+AA-Turner@users.noreply.github.com>
Co-authored-by: Hugo van Kemenade <1324225+hugovk@users.noreply.github.com>
Co-authored-by: Wingy <git@wingysam.xyz>
Co-authored-by: Koudai Aono <koxudaxi@gmail.com>
Co-authored-by: Dave Peck <davepeck@gmail.com>
Co-authored-by: Terry Jan Reedy <tjreedy@udel.edu>
Co-authored-by: Paul Everitt <pauleveritt@me.com>
Co-authored-by: sobolevn <mail@sobolevn.me>

1 files changed, 173 insertions, 28 deletions

diff --git a/Python/ast_unparse.c b/Python/ast_unparse.c
index f3c669c33eb..c121ec096ae 100644
--- a/Python/ast_unparse.c
+++ b/Python/ast_unparse.c
@@ -18,8 +18,12 @@ expr_as_unicode(expr_ty e, int level);
 static int
 append_ast_expr(PyUnicodeWriter *writer, expr_ty e, int level);
 static int
+append_templatestr(PyUnicodeWriter *writer, expr_ty e);
+static int
 append_joinedstr(PyUnicodeWriter *writer, expr_ty e, bool is_format_spec);
 static int
+append_interpolation(PyUnicodeWriter *writer, expr_ty e);
+static int
 append_formattedvalue(PyUnicodeWriter *writer, expr_ty e);
 static int
 append_ast_slice(PyUnicodeWriter *writer, expr_ty e);
@@ -621,11 +625,15 @@ append_fstring_element(PyUnicodeWriter *writer, expr_ty e, bool is_format_spec)
         return append_fstring_unicode(writer, e->v.Constant.value);
     case JoinedStr_kind:
         return append_joinedstr(writer, e, is_format_spec);
+    case TemplateStr_kind:
+        return append_templatestr(writer, e);
     case FormattedValue_kind:
         return append_formattedvalue(writer, e);
+    case Interpolation_kind:
+        return append_interpolation(writer, e);
     default:
         PyErr_SetString(PyExc_SystemError,
-                        "unknown expression kind inside f-string");
+                        "unknown expression kind inside f-string or t-string");
         return -1;
     }
 }
@@ -633,7 +641,7 @@ append_fstring_element(PyUnicodeWriter *writer, expr_ty e, bool is_format_spec)
 /* Build body separately to enable wrapping the entire stream of Strs,
    Constants and FormattedValues in one opening and one closing quote. */
 static PyObject *
-build_fstring_body(asdl_expr_seq *values, bool is_format_spec)
+build_ftstring_body(asdl_expr_seq *values, bool is_format_spec)
 {
     PyUnicodeWriter *body_writer = PyUnicodeWriter_Create(256);
     if (body_writer == NULL) {
@@ -655,10 +663,98 @@ build_fstring_body(asdl_expr_seq *values, bool is_format_spec)
 }
 
 static int
+_write_values_subarray(PyUnicodeWriter *writer, asdl_expr_seq *values, Py_ssize_t first_idx,
+                       Py_ssize_t last_idx, char prefix, PyArena *arena)
+{
+    int result = -1;
+
+    asdl_expr_seq *new_values = _Py_asdl_expr_seq_new(last_idx - first_idx + 1, arena);
+    if (!new_values) {
+        return result;
+    }
+
+    Py_ssize_t j = 0;
+    for (Py_ssize_t i = first_idx; i <= last_idx; ++i) {
+        asdl_seq_SET(new_values, j++, asdl_seq_GET(values, i));
+    }
+
+    PyObject *body = build_ftstring_body(new_values, false);
+    if (!body) {
+        return result;
+    }
+
+    if (-1 != append_char(writer, prefix) &&
+        -1 != append_repr(writer, body))
+    {
+        result = 0;
+    }
+    Py_DECREF(body);
+    return result;
+}
+
+static int
+append_templatestr(PyUnicodeWriter *writer, expr_ty e)
+{
+    PyArena *arena = _PyArena_New();
+    if (!arena) {
+        return -1;
+    }
+
+    Py_ssize_t last_idx = 0;
+    Py_ssize_t len = asdl_seq_LEN(e->v.TemplateStr.values);
+    for (Py_ssize_t i = 0; i < len; i++) {
+        expr_ty value = asdl_seq_GET(e->v.TemplateStr.values, i);
+
+        // Handle implicit concat of t-strings with f-strings
+        if (value->kind == FormattedValue_kind) {
+            if (i > last_idx) {
+                // Create a new TemplateStr with the values between last_idx and i
+                // and append it to the writer.
+                if (_write_values_subarray(writer, e->v.TemplateStr.values,
+                        last_idx, i - 1, 't', arena) == -1) {
+                    goto error;
+                }
+
+                if (append_charp(writer, " ") == -1) {
+                    goto error;
+                }
+            }
+
+            // Append the FormattedValue to the writer.
+            if (_write_values_subarray(writer, e->v.TemplateStr.values,
+                    i, i, 'f', arena) == -1) {
+                goto error;
+            }
+
+            if (i + 1 < len) {
+                if (append_charp(writer, " ") == -1) {
+                    goto error;
+                }
+            }
+
+            last_idx = i + 1;
+        }
+    }
+
+    if (last_idx < len) {
+        if (_write_values_subarray(writer, e->v.TemplateStr.values,
+                last_idx, len - 1, 't', arena) == -1) {
+            goto error;
+        }
+    }
+
+    return 0;
+
+error:
+    _PyArena_Free(arena);
+    return -1;
+}
+
+static int
 append_joinedstr(PyUnicodeWriter *writer, expr_ty e, bool is_format_spec)
 {
     int result = -1;
-    PyObject *body = build_fstring_body(e->v.JoinedStr.values, is_format_spec);
+    PyObject *body = build_ftstring_body(e->v.JoinedStr.values, is_format_spec);
     if (!body) {
         return -1;
     }
@@ -678,13 +774,12 @@ append_joinedstr(PyUnicodeWriter *writer, expr_ty e, bool is_format_spec)
 }
 
 static int
-append_formattedvalue(PyUnicodeWriter *writer, expr_ty e)
+append_interpolation_value(PyUnicodeWriter *writer, expr_ty e)
 {
-    const char *conversion;
     const char *outer_brace = "{";
     /* Grammar allows PR_TUPLE, but use >PR_TEST for adding parenthesis
        around a lambda with ':' */
-    PyObject *temp_fv_str = expr_as_unicode(e->v.FormattedValue.value, PR_TEST + 1);
+    PyObject *temp_fv_str = expr_as_unicode(e, PR_TEST + 1);
     if (!temp_fv_str) {
         return -1;
     }
@@ -702,35 +797,81 @@ append_formattedvalue(PyUnicodeWriter *writer, expr_ty e)
         return -1;
     }
     Py_DECREF(temp_fv_str);
+    return 0;
+}
 
-    if (e->v.FormattedValue.conversion > 0) {
-        switch (e->v.FormattedValue.conversion) {
-        case 'a':
-            conversion = "!a";
-            break;
-        case 'r':
-            conversion = "!r";
-            break;
-        case 's':
-            conversion = "!s";
-            break;
-        default:
-            PyErr_SetString(PyExc_SystemError,
-                            "unknown f-value conversion kind");
-            return -1;
-        }
-        APPEND_STR(conversion);
+static int
+append_interpolation_conversion(PyUnicodeWriter *writer, int conversion)
+{
+    if (conversion < 0) {
+        return 0;
     }
-    if (e->v.FormattedValue.format_spec) {
+
+    const char *conversion_str;
+    switch (conversion) {
+    case 'a':
+        conversion_str = "!a";
+        break;
+    case 'r':
+        conversion_str = "!r";
+        break;
+    case 's':
+        conversion_str = "!s";
+        break;
+    default:
+        PyErr_SetString(PyExc_SystemError,
+                        "unknown f-value conversion kind");
+        return -1;
+    }
+    APPEND_STR(conversion_str);
+    return 0;
+}
+
+static int
+append_interpolation_format_spec(PyUnicodeWriter *writer, expr_ty e)
+{
+    if (e) {
         if (-1 == PyUnicodeWriter_WriteChar(writer, ':') ||
-            -1 == append_fstring_element(writer,
-                                         e->v.FormattedValue.format_spec,
-                                         true
-                                        ))
+            -1 == append_fstring_element(writer, e, true))
         {
             return -1;
         }
     }
+    return 0;
+}
+
+static int
+append_interpolation(PyUnicodeWriter *writer, expr_ty e)
+{
+    if (-1 == append_interpolation_value(writer, e->v.Interpolation.value)) {
+        return -1;
+    }
+
+    if (-1 == append_interpolation_conversion(writer, e->v.Interpolation.conversion)) {
+        return -1;
+    }
+
+    if (-1 == append_interpolation_format_spec(writer, e->v.Interpolation.format_spec)) {
+        return -1;
+    }
+
+    APPEND_STR_FINISH("}");
+}
+
+static int
+append_formattedvalue(PyUnicodeWriter *writer, expr_ty e)
+{
+    if (-1 == append_interpolation_value(writer, e->v.FormattedValue.value)) {
+        return -1;
+    }
+
+    if (-1 == append_interpolation_conversion(writer, e->v.FormattedValue.conversion)) {
+        return -1;
+    }
+
+    if (-1 == append_interpolation_format_spec(writer, e->v.FormattedValue.format_spec)) {
+        return -1;
+    }
 
     APPEND_CHAR_FINISH('}');
 }
@@ -901,8 +1042,12 @@ append_ast_expr(PyUnicodeWriter *writer, expr_ty e, int level)
         return append_ast_constant(writer, e->v.Constant.value);
     case JoinedStr_kind:
         return append_joinedstr(writer, e, false);
+    case TemplateStr_kind:
+        return append_templatestr(writer, e);
     case FormattedValue_kind:
         return append_formattedvalue(writer, e);
+    case Interpolation_kind:
+        return append_interpolation(writer, e);
     /* The following exprs can be assignment targets. */
     case Attribute_kind:
         return append_ast_attribute(writer, e);