Skip to content

Commit bfac7d2

Browse files
gh-133581: Improve AST unparsing of t-strings (#133635)
1 parent a2c4467 commit bfac7d2

File tree

5 files changed

+52
-17
lines changed

5 files changed

+52
-17
lines changed

Lib/_ast_unparse.py

Lines changed: 9 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -627,6 +627,9 @@ def _write_ftstring(self, values, prefix):
627627
self._ftstring_helper(fstring_parts)
628628

629629
def _tstring_helper(self, node):
630+
if not node.values:
631+
self._write_ftstring([], "t")
632+
return
630633
last_idx = 0
631634
for i, value in enumerate(node.values):
632635
# This can happen if we have an implicit concat of a t-string
@@ -679,9 +682,12 @@ def _unparse_interpolation_value(self, inner):
679682
unparser.set_precedence(_Precedence.TEST.next(), inner)
680683
return unparser.visit(inner)
681684

682-
def _write_interpolation(self, node):
685+
def _write_interpolation(self, node, is_interpolation=False):
683686
with self.delimit("{", "}"):
684-
expr = self._unparse_interpolation_value(node.value)
687+
if is_interpolation:
688+
expr = node.str
689+
else:
690+
expr = self._unparse_interpolation_value(node.value)
685691
if expr.startswith("{"):
686692
# Separate pair of opening brackets as "{ {"
687693
self.write(" ")
@@ -696,7 +702,7 @@ def visit_FormattedValue(self, node):
696702
self._write_interpolation(node)
697703

698704
def visit_Interpolation(self, node):
699-
self._write_interpolation(node)
705+
self._write_interpolation(node, is_interpolation=True)
700706

701707
def visit_Name(self, node):
702708
self.write(node.id)

Lib/test/test_future_stmt/test_future.py

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -422,6 +422,11 @@ def test_annotations(self):
422422
eq('(((a)))', 'a')
423423
eq('(((a, b)))', '(a, b)')
424424
eq("1 + 2 + 3")
425+
eq("t''")
426+
eq("t'{a + b}'")
427+
eq("t'{a!s}'")
428+
eq("t'{a:b}'")
429+
eq("t'{a:b=}'")
425430

426431
def test_fstring_debug_annotations(self):
427432
# f-strings with '=' don't round trip very well, so set the expected

Lib/test/test_unparse.py

Lines changed: 9 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -817,6 +817,15 @@ def test_type_params(self):
817817
self.check_ast_roundtrip("def f[T: int = int, **P = int, *Ts = *int]():\n pass")
818818
self.check_ast_roundtrip("class C[T: int = int, **P = int, *Ts = *int]():\n pass")
819819

820+
def test_tstr(self):
821+
self.check_ast_roundtrip("t'{a + b}'")
822+
self.check_ast_roundtrip("t'{a + b:x}'")
823+
self.check_ast_roundtrip("t'{a + b!s}'")
824+
self.check_ast_roundtrip("t'{ {a}}'")
825+
self.check_ast_roundtrip("t'{ {a}=}'")
826+
self.check_ast_roundtrip("t'{{a}}'")
827+
self.check_ast_roundtrip("t''")
828+
820829

821830
class ManualASTCreationTestCase(unittest.TestCase):
822831
"""Test that AST nodes created without a type_params field unparse correctly."""
@@ -942,7 +951,6 @@ def files_to_test(cls):
942951
for directory in cls.test_directories
943952
for item in directory.glob("*.py")
944953
if not item.name.startswith("bad")
945-
and item.name != "annotationlib.py" # gh-133581: t"" does not roundtrip
946954
]
947955

948956
# Test limited subset of files unless the 'cpu' resource is specified.
Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,4 @@
1+
Improve unparsing of t-strings in :func:`ast.unparse` and ``from __future__
2+
import annotations``. Empty t-strings now round-trip correctly and
3+
formatting in interpolations is preserved.
4+
Patch by Jelle Zijlstra.

Python/ast_unparse.c

Lines changed: 25 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -702,6 +702,13 @@ append_templatestr(PyUnicodeWriter *writer, expr_ty e)
702702

703703
Py_ssize_t last_idx = 0;
704704
Py_ssize_t len = asdl_seq_LEN(e->v.TemplateStr.values);
705+
if (len == 0) {
706+
int result = _write_values_subarray(writer, e->v.TemplateStr.values,
707+
0, len - 1, 't', arena);
708+
_PyArena_Free(arena);
709+
return result;
710+
}
711+
705712
for (Py_ssize_t i = 0; i < len; i++) {
706713
expr_ty value = asdl_seq_GET(e->v.TemplateStr.values, i);
707714

@@ -774,32 +781,37 @@ append_joinedstr(PyUnicodeWriter *writer, expr_ty e, bool is_format_spec)
774781
}
775782

776783
static int
777-
append_interpolation_value(PyUnicodeWriter *writer, expr_ty e)
784+
append_interpolation_str(PyUnicodeWriter *writer, PyObject *str)
778785
{
779786
const char *outer_brace = "{";
780-
/* Grammar allows PR_TUPLE, but use >PR_TEST for adding parenthesis
781-
around a lambda with ':' */
782-
PyObject *temp_fv_str = expr_as_unicode(e, PR_TEST + 1);
783-
if (!temp_fv_str) {
784-
return -1;
785-
}
786-
if (PyUnicode_Find(temp_fv_str, _Py_LATIN1_CHR('{'), 0, 1, 1) == 0) {
787+
if (PyUnicode_Find(str, _Py_LATIN1_CHR('{'), 0, 1, 1) == 0) {
787788
/* Expression starts with a brace, split it with a space from the outer
788789
one. */
789790
outer_brace = "{ ";
790791
}
791792
if (-1 == append_charp(writer, outer_brace)) {
792-
Py_DECREF(temp_fv_str);
793793
return -1;
794794
}
795-
if (-1 == PyUnicodeWriter_WriteStr(writer, temp_fv_str)) {
796-
Py_DECREF(temp_fv_str);
795+
if (-1 == PyUnicodeWriter_WriteStr(writer, str)) {
797796
return -1;
798797
}
799-
Py_DECREF(temp_fv_str);
800798
return 0;
801799
}
802800

801+
static int
802+
append_interpolation_value(PyUnicodeWriter *writer, expr_ty e)
803+
{
804+
/* Grammar allows PR_TUPLE, but use >PR_TEST for adding parenthesis
805+
around a lambda with ':' */
806+
PyObject *temp_fv_str = expr_as_unicode(e, PR_TEST + 1);
807+
if (!temp_fv_str) {
808+
return -1;
809+
}
810+
int result = append_interpolation_str(writer, temp_fv_str);
811+
Py_DECREF(temp_fv_str);
812+
return result;
813+
}
814+
803815
static int
804816
append_interpolation_conversion(PyUnicodeWriter *writer, int conversion)
805817
{
@@ -843,7 +855,7 @@ append_interpolation_format_spec(PyUnicodeWriter *writer, expr_ty e)
843855
static int
844856
append_interpolation(PyUnicodeWriter *writer, expr_ty e)
845857
{
846-
if (-1 == append_interpolation_value(writer, e->v.Interpolation.value)) {
858+
if (-1 == append_interpolation_str(writer, e->v.Interpolation.str)) {
847859
return -1;
848860
}
849861

0 commit comments

Comments
 (0)