Fix unparsing implicit concat of t-strings with f-strings

Author: lysnikolaou

Lib/_ast_unparse.py

@@ -573,7 +573,7 @@ def _write_str_avoiding_backslashes(self, string, *, quote_types=_ALL_QUOTES):
         quote_type = quote_types[0]
         self.write(f"{quote_type}{string}{quote_type}")
 
-    def _ftstring_helper(self, node, ftstring_parts):
+    def _ftstring_helper(self, ftstring_parts):
         new_ftstring_parts = []
         quote_types = list(_ALL_QUOTES)
         fallback_to_repr = False
@@ -615,22 +615,43 @@ def _ftstring_helper(self, node, ftstring_parts):
         quote_type = quote_types[0]
         self.write(f"{quote_type}{value}{quote_type}")
 
-    def _write_ftstring(self, node, prefix):
+    def _write_ftstring(self, values, prefix):
         self.write(prefix)
         fstring_parts = []
-        for value in node.values:
+        for value in values:
             with self.buffered() as buffer:
                 self._write_ftstring_inner(value)
             fstring_parts.append(
                 ("".join(buffer), isinstance(value, Constant))
             )
-        self._ftstring_helper(node, fstring_parts)
+        self._ftstring_helper(fstring_parts)
+
+    def _tstring_helper(self, node):
+        last_idx = 0
+        for i, value in enumerate(node.values):
+            # This can happen if we have an implicit concat of a t-string
+            # with an f-string
+            if isinstance(value, FormattedValue):
+                if i > last_idx:
+                    # Write t-string until here
+                    self._write_ftstring(node.values[last_idx:i], "t")
+                    self.write(" ")
+                # Write f-string with the current formatted value
+                self._write_ftstring([node.values[i]], "f")
+                if i + 1 < len(node.values):
+                    # Only add a space if there are more values after this
+                    self.write(" ")
+                last_idx = i + 1
+
+        if last_idx < len(node.values):
+            # Write t-string from last_idx to end
+            self._write_ftstring(node.values[last_idx:], "t")
 
     def visit_JoinedStr(self, node):
-        self._write_ftstring(node, "f")
+        self._write_ftstring(node.values, "f")
 
     def visit_TemplateStr(self, node):
-        self._write_ftstring(node, "t")
+        self._tstring_helper(node)
 
     def _write_ftstring_inner(self, node, is_format_spec=False):
         if isinstance(node, JoinedStr):

Lib/test/support/ast_helper.py

@@ -16,6 +16,9 @@ def traverse_compare(a, b, missing=object()):
                 self.fail(f"{type(a)!r} is not {type(b)!r}")
             if isinstance(a, ast.AST):
                 for field in a._fields:
+                    if isinstance(a, ast.Constant) and field == "kind":
+                        # Skip the 'kind' field for ast.Constant
+                        continue
                     value1 = getattr(a, field, missing)
                     value2 = getattr(b, field, missing)
                     # Singletons are equal by definition, so further

Lib/test/test_unparse.py

@@ -202,6 +202,15 @@ def test_fstrings_pep701(self):
         self.check_ast_roundtrip('f" something { my_dict["key"] } something else "')
         self.check_ast_roundtrip('f"{f"{f"{f"{f"{f"{1+1}"}"}"}"}"}"')
 
+    def test_tstrings(self):
+        self.check_ast_roundtrip("t'foo'")
+        self.check_ast_roundtrip("t'foo {bar}'")
+        self.check_ast_roundtrip("t'foo {bar!s:.2f}'")
+        self.check_ast_roundtrip("t'foo {bar}' f'{bar}'")
+        self.check_ast_roundtrip("f'{bar}' t'foo {bar}'")
+        self.check_ast_roundtrip("t'foo {bar}' fr'\hello {bar}'")
+        self.check_ast_roundtrip("t'foo {bar}' u'bar'")
+
     def test_strings(self):
         self.check_ast_roundtrip("u'foo'")
         self.check_ast_roundtrip("r'foo'")

Python/ast_unparse.c

@@ -663,15 +663,27 @@ build_ftstring_body(asdl_expr_seq *values, bool is_format_spec)
 }
 
 static int
-append_templatestr(PyUnicodeWriter *writer, expr_ty e)
+_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;
-    PyObject *body = build_ftstring_body(e->v.TemplateStr.values, 0);
+
+    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_charp(writer, "t") &&
+    if (-1 != append_char(writer, prefix) &&
         -1 != append_repr(writer, body))
     {
         result = 0;
@@ -680,6 +692,64 @@ append_templatestr(PyUnicodeWriter *writer, expr_ty e)
     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)
 {