Fix test_unparse again!

Author: lysnikolaou

Lib/test/test_tstring.py

@@ -207,7 +207,7 @@ def test_debug_specifier(self):
         t = t"Value: {value=}"
         self.assertEqual(t.strings, ("Value: value=", ""))
         self.assertEqual(t.interpolations[0].value, value)
-        self.assertEqual(t.interpolations[0].expression, "value=")
+        self.assertEqual(t.interpolations[0].expression, "value")
         self.assertEqual(t.interpolations[0].conversion, "r")
         self.assertEqual(t.interpolations[0].format_spec, "")
         self.assertEqual(f(t), "Value: value=42")
@@ -216,7 +216,7 @@ def test_debug_specifier(self):
         t = t"Value: {value=:.2f}"
         self.assertEqual(t.strings, ("Value: value=", ""))
         self.assertEqual(t.interpolations[0].value, value)
-        self.assertEqual(t.interpolations[0].expression, "value=")
+        self.assertEqual(t.interpolations[0].expression, "value")
         self.assertEqual(t.interpolations[0].conversion, None)
         self.assertEqual(t.interpolations[0].format_spec, ".2f")
         self.assertEqual(f(t), "Value: value=42.00")
@@ -225,15 +225,15 @@ def test_debug_specifier(self):
         t = t"Value: {value=!s}"
         self.assertEqual(t.strings, ("Value: value=", ""))
         self.assertEqual(t.interpolations[0].value, value)
-        self.assertEqual(t.interpolations[0].expression, "value=")
+        self.assertEqual(t.interpolations[0].expression, "value")
         self.assertEqual(t.interpolations[0].conversion, "s")
         self.assertEqual(t.interpolations[0].format_spec, "")
 
         # Test white space in debug specifier
         t = t"Value: {value = }"
         self.assertEqual(t.strings, ("Value: value = ", ""))
         self.assertEqual(t.interpolations[0].value, value)
-        self.assertEqual(t.interpolations[0].expression, "value = ")
+        self.assertEqual(t.interpolations[0].expression, "value")
         self.assertEqual(t.interpolations[0].conversion, "r")
         self.assertEqual(t.interpolations[0].format_spec, "")
         self.assertEqual(f(t), "Value: value = 42")

Lib/test/test_unparse.py

@@ -918,7 +918,7 @@ class DirectoryTestCase(ASTTestCase):
     run_always_files = {"test_grammar.py", "test_syntax.py", "test_compile.py",
                         "test_ast.py", "test_asdl_parser.py", "test_fstring.py",
                         "test_patma.py", "test_type_alias.py", "test_type_params.py",
-                        "test_tokenize.py"}
+                        "test_tokenize.py", "test_tstring.py"}
 
     _files_to_test = None
 

Parser/action_helpers.c

@@ -1475,6 +1475,24 @@ _get_interpolation_conversion(Parser *p, Token *debug, ResultTokenWithMetadata *
     return -1;
 }
 
+static PyObject *
+_strip_interpolation_expr(PyObject *exprstr)
+{
+    Py_ssize_t len = PyUnicode_GET_LENGTH(exprstr);
+
+    for (Py_ssize_t i = len - 1; i >= 0; i--) {
+        Py_UCS4 c = PyUnicode_READ_CHAR(exprstr, i);
+        if (_PyUnicode_IsWhitespace(c) || c == '=') {
+            len--;
+        }
+        else {
+            break;
+        }
+    }
+
+    return PyUnicode_Substring(exprstr, 0, len);
+}
+
 expr_ty _PyPegen_interpolation(Parser *p, expr_ty expression, Token *debug, ResultTokenWithMetadata *conversion,
                                  ResultTokenWithMetadata *format, Token *closing_brace, int lineno, int col_offset,
                                  int end_lineno, int end_col_offset, PyArena *arena) {
@@ -1485,7 +1503,8 @@ expr_ty _PyPegen_interpolation(Parser *p, expr_ty expression, Token *debug, Resu
         char buf[1];
         buf[0] = conversion_val;
         convstr = PyUnicode_FromStringAndSize(buf, 1);
-        if (convstr == NULL) {
+        if (convstr == NULL || _PyArena_AddPyObject(arena, convstr) < 0) {
+            Py_XDECREF(convstr);
             return NULL;
         }
     }
@@ -1512,8 +1531,14 @@ expr_ty _PyPegen_interpolation(Parser *p, expr_ty expression, Token *debug, Resu
     }
 
     assert(exprstr != NULL);
+    PyObject *final_exprstr = _strip_interpolation_expr(exprstr);
+    if (!final_exprstr || _PyArena_AddPyObject(arena, final_exprstr) < 0) {
+        Py_XDECREF(final_exprstr);
+        return NULL;
+    }
+
     expr_ty interpolation = _PyAST_Interpolation(
-        expression, exprstr, convstr, format ? (expr_ty) format->result : NULL,
+        expression, final_exprstr, convstr, format ? (expr_ty) format->result : NULL,
         lineno, col_offset, end_lineno,
         end_col_offset, arena
     );