Use value field to generate code when None in str field of ast.Interpolation

Author: George-Ogden

Lib/_ast_unparse.py

@@ -658,9 +658,9 @@ def _unparse_interpolation_value(self, inner):
         unparser.set_precedence(_Precedence.TEST.next(), inner)
         return unparser.visit(inner)
 
-    def _write_interpolation(self, node, is_interpolation=False):
+    def _write_interpolation(self, node, use_str_attr=False):
         with self.delimit("{", "}"):
-            if is_interpolation:
+            if use_str_attr:
                 expr = node.str
             else:
                 expr = self._unparse_interpolation_value(node.value)
@@ -678,7 +678,8 @@ def visit_FormattedValue(self, node):
         self._write_interpolation(node)
 
     def visit_Interpolation(self, node):
-        self._write_interpolation(node, is_interpolation=True)
+        # If `str` is set to `None`, use the `value` to generate the source code.
+        self._write_interpolation(node, use_str_attr=node.str is not None)
 
     def visit_Name(self, node):
         self.write(node.id)

Lib/test/test_unparse.py

@@ -215,6 +215,88 @@ def test_tstrings(self):
         self.check_ast_roundtrip("t''")
         self.check_ast_roundtrip('t""')
         self.check_ast_roundtrip("t'{(lambda x: x)}'")
+        self.check_ast_roundtrip("t'{t'{x}'}'")
+
+    def test_tstring_with_nonsensical_str_field(self):
+        # `value` suggests that the original code is `t'{test1}`, but `str` suggests otherwise
+        self.assertEqual(
+            ast.unparse(
+                ast.TemplateStr(
+                    values=[
+                        ast.Interpolation(
+                            value=ast.Name(id="test1", ctx=ast.Load()), str="test2", conversion=-1
+                        )
+                    ]
+                )
+            ),
+            "t'{test2}'",
+        )
+
+    def test_tstring_with_none_str_field(self):
+        self.assertEqual(
+            ast.unparse(
+                ast.TemplateStr(
+                    [ast.Interpolation(value=ast.Name(id="test1"), str=None, conversion=-1)]
+                )
+            ),
+            "t'{test1}'",
+        )
+        self.assertEqual(
+            ast.unparse(
+                ast.TemplateStr(
+                    [
+                        ast.Interpolation(
+                            value=ast.Lambda(
+                                args=ast.arguments(args=[ast.arg(arg="x")]),
+                                body=ast.Name(id="x"),
+                            ),
+                            str=None,
+                            conversion=-1,
+                        )
+                    ]
+                )
+            ),
+            "t'{(lambda x: x)}'",
+        )
+        self.assertEqual(
+            ast.unparse(
+                ast.TemplateStr(
+                    values=[
+                        ast.Interpolation(
+                            value=ast.TemplateStr(
+                                # `str` field kept here
+                                [ast.Interpolation(value=ast.Name(id="x"), str="y", conversion=-1)]
+                            ),
+                            str=None,
+                            conversion=-1,
+                        )
+                    ]
+                )
+            ),
+            '''t"{t'{y}'}"''',
+        )
+        self.assertEqual(
+            ast.unparse(
+                ast.TemplateStr(
+                    values=[
+                        ast.Interpolation(
+                            value=ast.TemplateStr(
+                                [ast.Interpolation(value=ast.Name(id="x"), str=None, conversion=-1)]
+                            ),
+                            str=None,
+                            conversion=-1,
+                        )
+                    ]
+                )
+            ),
+            '''t"{t'{x}'}"''',
+        )
+        self.assertEqual(
+            ast.unparse(ast.TemplateStr(
+                [ast.Interpolation(value=ast.Constant(value="foo"), str=None, conversion=114)]
+            )),
+            '''t"{'foo'!r}"''',
+        )
 
     def test_strings(self):
         self.check_ast_roundtrip("u'foo'")