gh-133581: Improve AST unparsing of t-strings (#133635)

This commit is contained in:
Jelle Zijlstra 2025-05-08 06:13:57 -07:00 committed by GitHub
parent a2c4467d06
commit bfac7d2edc
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
5 changed files with 55 additions and 20 deletions

View file

@ -627,6 +627,9 @@ class Unparser(NodeVisitor):
self._ftstring_helper(fstring_parts)
def _tstring_helper(self, node):
if not node.values:
self._write_ftstring([], "t")
return
last_idx = 0
for i, value in enumerate(node.values):
# This can happen if we have an implicit concat of a t-string
@ -679,9 +682,12 @@ class Unparser(NodeVisitor):
unparser.set_precedence(_Precedence.TEST.next(), inner)
return unparser.visit(inner)
def _write_interpolation(self, node):
def _write_interpolation(self, node, is_interpolation=False):
with self.delimit("{", "}"):
expr = self._unparse_interpolation_value(node.value)
if is_interpolation:
expr = node.str
else:
expr = self._unparse_interpolation_value(node.value)
if expr.startswith("{"):
# Separate pair of opening brackets as "{ {"
self.write(" ")
@ -696,7 +702,7 @@ class Unparser(NodeVisitor):
self._write_interpolation(node)
def visit_Interpolation(self, node):
self._write_interpolation(node)
self._write_interpolation(node, is_interpolation=True)
def visit_Name(self, node):
self.write(node.id)

View file

@ -422,6 +422,11 @@ class AnnotationsFutureTestCase(unittest.TestCase):
eq('(((a)))', 'a')
eq('(((a, b)))', '(a, b)')
eq("1 + 2 + 3")
eq("t''")
eq("t'{a + b}'")
eq("t'{a!s}'")
eq("t'{a:b}'")
eq("t'{a:b=}'")
def test_fstring_debug_annotations(self):
# f-strings with '=' don't round trip very well, so set the expected

View file

@ -817,6 +817,15 @@ class CosmeticTestCase(ASTTestCase):
self.check_ast_roundtrip("def f[T: int = int, **P = int, *Ts = *int]():\n pass")
self.check_ast_roundtrip("class C[T: int = int, **P = int, *Ts = *int]():\n pass")
def test_tstr(self):
self.check_ast_roundtrip("t'{a + b}'")
self.check_ast_roundtrip("t'{a + b:x}'")
self.check_ast_roundtrip("t'{a + b!s}'")
self.check_ast_roundtrip("t'{ {a}}'")
self.check_ast_roundtrip("t'{ {a}=}'")
self.check_ast_roundtrip("t'{{a}}'")
self.check_ast_roundtrip("t''")
class ManualASTCreationTestCase(unittest.TestCase):
"""Test that AST nodes created without a type_params field unparse correctly."""
@ -942,7 +951,6 @@ class DirectoryTestCase(ASTTestCase):
for directory in cls.test_directories
for item in directory.glob("*.py")
if not item.name.startswith("bad")
and item.name != "annotationlib.py" # gh-133581: t"" does not roundtrip
]
# Test limited subset of files unless the 'cpu' resource is specified.

View file

@ -0,0 +1,4 @@
Improve unparsing of t-strings in :func:`ast.unparse` and ``from __future__
import annotations``. Empty t-strings now round-trip correctly and
formatting in interpolations is preserved.
Patch by Jelle Zijlstra.

View file

@ -702,6 +702,13 @@ append_templatestr(PyUnicodeWriter *writer, expr_ty e)
Py_ssize_t last_idx = 0;
Py_ssize_t len = asdl_seq_LEN(e->v.TemplateStr.values);
if (len == 0) {
int result = _write_values_subarray(writer, e->v.TemplateStr.values,
0, len - 1, 't', arena);
_PyArena_Free(arena);
return result;
}
for (Py_ssize_t i = 0; i < len; i++) {
expr_ty value = asdl_seq_GET(e->v.TemplateStr.values, i);
@ -774,30 +781,35 @@ append_joinedstr(PyUnicodeWriter *writer, expr_ty e, bool is_format_spec)
}
static int
append_interpolation_value(PyUnicodeWriter *writer, expr_ty e)
append_interpolation_str(PyUnicodeWriter *writer, PyObject *str)
{
const char *outer_brace = "{";
if (PyUnicode_Find(str, _Py_LATIN1_CHR('{'), 0, 1, 1) == 0) {
/* Expression starts with a brace, split it with a space from the outer
one. */
outer_brace = "{ ";
}
if (-1 == append_charp(writer, outer_brace)) {
return -1;
}
if (-1 == PyUnicodeWriter_WriteStr(writer, str)) {
return -1;
}
return 0;
}
static int
append_interpolation_value(PyUnicodeWriter *writer, expr_ty e)
{
/* Grammar allows PR_TUPLE, but use >PR_TEST for adding parenthesis
around a lambda with ':' */
PyObject *temp_fv_str = expr_as_unicode(e, PR_TEST + 1);
if (!temp_fv_str) {
return -1;
}
if (PyUnicode_Find(temp_fv_str, _Py_LATIN1_CHR('{'), 0, 1, 1) == 0) {
/* Expression starts with a brace, split it with a space from the outer
one. */
outer_brace = "{ ";
}
if (-1 == append_charp(writer, outer_brace)) {
Py_DECREF(temp_fv_str);
return -1;
}
if (-1 == PyUnicodeWriter_WriteStr(writer, temp_fv_str)) {
Py_DECREF(temp_fv_str);
return -1;
}
int result = append_interpolation_str(writer, temp_fv_str);
Py_DECREF(temp_fv_str);
return 0;
return result;
}
static int
@ -843,7 +855,7 @@ append_interpolation_format_spec(PyUnicodeWriter *writer, expr_ty e)
static int
append_interpolation(PyUnicodeWriter *writer, expr_ty e)
{
if (-1 == append_interpolation_value(writer, e->v.Interpolation.value)) {
if (-1 == append_interpolation_str(writer, e->v.Interpolation.str)) {
return -1;
}