Fix parser bug and implement ast_unparse in Python

Author: lysnikolaou

Lib/_ast_unparse.py

@@ -573,21 +573,11 @@ 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 visit_JoinedStr(self, node):
-        self.write("f")
-
-        fstring_parts = []
-        for value in node.values:
-            with self.buffered() as buffer:
-                self._write_fstring_inner(value)
-            fstring_parts.append(
-                ("".join(buffer), isinstance(value, Constant))
-            )
-
-        new_fstring_parts = []
+    def _ftstring_helper(self, node, ftstring_parts):
+        new_ftstring_parts = []
         quote_types = list(_ALL_QUOTES)
         fallback_to_repr = False
-        for value, is_constant in fstring_parts:
+        for value, is_constant in ftstring_parts:
             if is_constant:
                 value, new_quote_types = self._str_literal_helper(
                     value,
@@ -606,30 +596,47 @@ def visit_JoinedStr(self, node):
                 new_quote_types = [q for q in quote_types if q not in value]
                 if new_quote_types:
                     quote_types = new_quote_types
-            new_fstring_parts.append(value)
+            new_ftstring_parts.append(value)
 
         if fallback_to_repr:
             # If we weren't able to find a quote type that works for all parts
             # of the JoinedStr, fallback to using repr and triple single quotes.
             quote_types = ["'''"]
-            new_fstring_parts.clear()
-            for value, is_constant in fstring_parts:
+            new_ftstring_parts.clear()
+            for value, is_constant in ftstring_parts:
                 if is_constant:
                     value = repr('"' + value)  # force repr to use single quotes
                     expected_prefix = "'\""
                     assert value.startswith(expected_prefix), repr(value)
                     value = value[len(expected_prefix):-1]
-                new_fstring_parts.append(value)
+                new_ftstring_parts.append(value)
 
-        value = "".join(new_fstring_parts)
+        value = "".join(new_ftstring_parts)
         quote_type = quote_types[0]
         self.write(f"{quote_type}{value}{quote_type}")
 
-    def _write_fstring_inner(self, node, is_format_spec=False):
+    def _write_ftstring(self, node, prefix):
+        self.write(prefix)
+        fstring_parts = []
+        for value in node.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)
+
+    def visit_JoinedStr(self, node):
+        self._write_ftstring(node, "f")
+
+    def visit_TemplateStr(self, node):
+        self._write_ftstring(node, "t")
+
+    def _write_ftstring_inner(self, node, is_format_spec=False):
         if isinstance(node, JoinedStr):
             # for both the f-string itself, and format_spec
             for value in node.values:
-                self._write_fstring_inner(value, is_format_spec=is_format_spec)
+                self._write_ftstring_inner(value, is_format_spec=is_format_spec)
         elif isinstance(node, Constant) and isinstance(node.value, str):
             value = node.value.replace("{", "{{").replace("}", "}}")
 
@@ -641,26 +648,41 @@ def _write_fstring_inner(self, node, is_format_spec=False):
             self.write(value)
         elif isinstance(node, FormattedValue):
             self.visit_FormattedValue(node)
+        elif isinstance(node, Interpolation):
+            self.visit_Interpolation(node)
         else:
             raise ValueError(f"Unexpected node inside JoinedStr, {node!r}")
 
-    def visit_FormattedValue(self, node):
-        def unparse_inner(inner):
-            unparser = type(self)()
-            unparser.set_precedence(_Precedence.TEST.next(), inner)
-            return unparser.visit(inner)
+    def _unparse_interpolation_value(self, inner):
+        unparser = type(self)()
+        unparser.set_precedence(_Precedence.TEST.next(), inner)
+        return unparser.visit(inner)
+
+    def _write_fstring_conversion(self, node):
+        if node.conversion != -1:
+            self.write(f"!{chr(node.conversion)}")
+
+    def _write_tstring_conversion(self, node):
+        if node.conversion is not None:
+            self.write(f"!{node.conversion}")
 
+    def _write_interpolation(self, node, write_conversion):
         with self.delimit("{", "}"):
-            expr = unparse_inner(node.value)
+            expr = self._unparse_interpolation_value(node.value)
             if expr.startswith("{"):
                 # Separate pair of opening brackets as "{ {"
                 self.write(" ")
             self.write(expr)
-            if node.conversion != -1:
-                self.write(f"!{chr(node.conversion)}")
+            write_conversion(node)
             if node.format_spec:
                 self.write(":")
-                self._write_fstring_inner(node.format_spec, is_format_spec=True)
+                self._write_ftstring_inner(node.format_spec, is_format_spec=True)
+
+    def visit_FormattedValue(self, node):
+        self._write_interpolation(node, self._write_fstring_conversion)
+
+    def visit_Interpolation(self, node):
+        self._write_interpolation(node, self._write_tstring_conversion)
 
     def visit_Name(self, node):
         self.write(node.id)

Parser/action_helpers.c

@@ -1654,8 +1654,8 @@ _build_concatenated_unicode(Parser *p, asdl_expr_seq *strings, int lineno,
                            end_lineno, end_col_offset, arena);
 }
 
-static expr_ty
-_build_concatenated_joined_str(Parser *p, asdl_expr_seq *strings,
+static asdl_expr_seq *
+_build_concatenated_str(Parser *p, asdl_expr_seq *strings,
                                int lineno, int col_offset, int end_lineno,
                                int end_col_offset, PyArena *arena)
 {
@@ -1669,6 +1669,9 @@ _build_concatenated_joined_str(Parser *p, asdl_expr_seq *strings,
             case JoinedStr_kind:
                 n_flattened_elements += asdl_seq_LEN(elem->v.JoinedStr.values);
                 break;
+            case TemplateStr_kind:
+                n_flattened_elements += asdl_seq_LEN(elem->v.TemplateStr.values);
+                break;
             default:
                 n_flattened_elements++;
                 break;
@@ -1695,6 +1698,15 @@ _build_concatenated_joined_str(Parser *p, asdl_expr_seq *strings,
                     asdl_seq_SET(flattened, current_pos++, subvalue);
                 }
                 break;
+            case TemplateStr_kind:
+                for (Py_ssize_t j = 0; j < asdl_seq_LEN(elem->v.TemplateStr.values); j++) {
+                    expr_ty subvalue = asdl_seq_GET(elem->v.TemplateStr.values, j);
+                    if (subvalue == NULL) {
+                        return NULL;
+                    }
+                    asdl_seq_SET(flattened, current_pos++, subvalue);
+                }
+                break;
             default:
                 asdl_seq_SET(flattened, current_pos++, elem);
                 break;
@@ -1795,6 +1807,16 @@ _build_concatenated_joined_str(Parser *p, asdl_expr_seq *strings,
     }
 
     assert(current_pos == n_elements);
+    return values;
+}
+
+static expr_ty
+_build_concatenated_joined_str(Parser *p, asdl_expr_seq *strings,
+                               int lineno, int col_offset, int end_lineno,
+                               int end_col_offset, PyArena *arena)
+{
+    asdl_expr_seq *values = _build_concatenated_str(p, strings, lineno,
+        col_offset, end_lineno, end_col_offset, arena);
     return _PyAST_JoinedStr(values, lineno, col_offset, end_lineno, end_col_offset, p->arena);
 }
 
@@ -1803,53 +1825,9 @@ _build_concatenated_template_str(Parser *p, asdl_expr_seq *strings,
                                int lineno, int col_offset, int end_lineno,
                                int end_col_offset, PyArena *arena)
 {
-    Py_ssize_t len = asdl_seq_LEN(strings);
-    assert(len > 0);
-
-    Py_ssize_t n_flattened_elements = 0;
-    for (Py_ssize_t i = 0; i < len; i++) {
-        expr_ty elem = asdl_seq_GET(strings, i);
-        switch(elem->kind) {
-            case TemplateStr_kind:
-                n_flattened_elements += asdl_seq_LEN(elem->v.JoinedStr.values);
-                break;
-            default:
-                n_flattened_elements++;
-                break;
-        }
-    }
-
-
-    asdl_expr_seq* flattened = _Py_asdl_expr_seq_new(n_flattened_elements, p->arena);
-    if (flattened == NULL) {
-        return NULL;
-    }
-
-    Py_ssize_t pos = 0;
-    for (Py_ssize_t i = 0; i < len; i++) {
-        expr_ty elem = asdl_seq_GET(strings, i);
-
-        switch (elem->kind) {
-            case TemplateStr_kind:
-                for (Py_ssize_t j = 0; j < asdl_seq_LEN(elem->v.TemplateStr.values); j++) {
-                    expr_ty subitem = asdl_seq_GET(elem->v.TemplateStr.values, j);
-                    asdl_seq_SET(flattened, pos++, subitem);
-                }
-                break;
-            case JoinedStr_kind: {
-                expr_ty joined_str = _build_concatenated_joined_str(p,
-                    elem->v.JoinedStr.values, lineno, col_offset,
-                    end_lineno, end_col_offset, arena);
-                asdl_seq_SET(flattened, pos++, joined_str);
-                break;
-            }
-            default:
-                asdl_seq_SET(flattened, pos++, elem);
-                break;
-        }
-    }
-
-    return _PyAST_TemplateStr(flattened, lineno, col_offset, end_lineno,
+    asdl_expr_seq *values = _build_concatenated_str(p, strings, lineno,
+        col_offset, end_lineno, end_col_offset, arena);
+    return _PyAST_TemplateStr(values, lineno, col_offset, end_lineno,
         end_col_offset, arena);
 }