pythongh-116022: Improve repr() of AST nodes (python#117046)

Co-authored-by: AN Long <[email protected]> Co-authored-by: Jelle Zijlstra <[email protected]> Co-authored-by: Alex Waygood <[email protected]> Co-authored-by: Bénédikt Tran <[email protected]>
colesbury · Sep 18, 2024 · 21d2a9a · 21d2a9a
1 parent f9fa6ba
commit 21d2a9a
Show file tree

Hide file tree

Showing 7 changed files with 682 additions and 2 deletions.
diff --git a/Doc/library/ast.rst b/Doc/library/ast.rst
@@ -134,6 +134,11 @@ Node classes
    Simple indices are represented by their value, extended slices are
    represented as tuples.
 
+.. versionchanged:: 3.14
+
+    The :meth:`~object.__repr__` output of :class:`~ast.AST` nodes includes
+    the values of the node fields.
+
 .. deprecated:: 3.8
 
    Old classes :class:`!ast.Num`, :class:`!ast.Str`, :class:`!ast.Bytes`,

diff --git a/Lib/test/test_ast/data/ast_repr.txt b/Lib/test/test_ast/data/ast_repr.txt
diff --git a/Lib/test/test_ast/test_ast.py b/Lib/test/test_ast/test_ast.py
@@ -10,6 +10,7 @@
 import types
 import unittest
 import weakref
+from pathlib import Path
 from textwrap import dedent
 try:
     import _testinternalcapi
@@ -29,6 +30,16 @@
 STDLIB_FILES = [fn for fn in os.listdir(STDLIB) if fn.endswith(".py")]
 STDLIB_FILES.extend(["test/test_grammar.py", "test/test_unpack_ex.py"])
 
+AST_REPR_DATA_FILE = Path(__file__).parent / "data" / "ast_repr.txt"
+
+def ast_repr_get_test_cases() -> list[str]:
+    return exec_tests + eval_tests
+
+
+def ast_repr_update_snapshots() -> None:
+    data = [repr(ast.parse(test)) for test in ast_repr_get_test_cases()]
+    AST_REPR_DATA_FILE.write_text("\n".join(data))
+
 
 class AST_Tests(unittest.TestCase):
     maxDiff = None
@@ -408,7 +419,7 @@ def test_invalid_sum(self):
         m = ast.Module([ast.Expr(ast.expr(**pos), **pos)], [])
         with self.assertRaises(TypeError) as cm:
             compile(m, "<test>", "exec")
-        self.assertIn("but got <ast.expr", str(cm.exception))
+        self.assertIn("but got expr()", str(cm.exception))
 
     def test_invalid_identifier(self):
         m = ast.Module([ast.Expr(ast.Name(42, ast.Load()))], [])
@@ -772,6 +783,12 @@ def test_none_checks(self) -> None:
         for node, attr, source in tests:
             self.assert_none_check(node, attr, source)
 
+    def test_repr(self) -> None:
+        snapshots = AST_REPR_DATA_FILE.read_text().split("\n")
+        for test, snapshot in zip(ast_repr_get_test_cases(), snapshots, strict=True):
+            with self.subTest(test_input=test):
+                self.assertEqual(repr(ast.parse(test)), snapshot)
+
 
 class CopyTests(unittest.TestCase):
     """Test copying and pickling AST nodes."""
@@ -3332,5 +3349,8 @@ def test_folding_type_param_in_type_alias(self):
             self.assert_ast(result_code, non_optimized_target, optimized_target)
 
 
-if __name__ == "__main__":
+if __name__ == '__main__':
+    if len(sys.argv) > 1 and sys.argv[1] == '--snapshot-update':
+        ast_repr_update_snapshots()
+        sys.exit(0)
     unittest.main()
diff --git a/Makefile.pre.in b/Makefile.pre.in
@@ -2416,6 +2416,7 @@ LIBSUBDIRS=	asyncio \
 TESTSUBDIRS=	idlelib/idle_test \
 		test \
 		test/test_ast \
+		test/test_ast/data \
 		test/archivetestdata \
 		test/audiodata \
 		test/certdata \

diff --git a/Misc/NEWS.d/next/Core and Builtins/2024-03-19-22-21-22.gh-issue-116022.iyHENN.rst b/Misc/NEWS.d/next/Core and Builtins/2024-03-19-22-21-22.gh-issue-116022.iyHENN.rst
@@ -0,0 +1 @@
+Improve the :meth:`~object.__repr__` output of :class:`~ast.AST` nodes.
diff --git a/Parser/asdl_c.py b/Parser/asdl_c.py
@@ -1435,8 +1435,230 @@ def visitModule(self, mod):
     {NULL}
 };
 
+static PyObject *
+ast_repr_max_depth(AST_object *self, int depth);
+
+/* Format list and tuple properties of AST nodes.
+   Note that, only the first and last elements are shown.
+   Anything in between is represented with an ellipsis ('...').
+   For example, the list [1, 2, 3] is formatted as
+   'List(elts=[Constant(1), ..., Constant(3)])'. */
+static PyObject *
+ast_repr_list(PyObject *list, int depth)
+{
+    assert(PyList_Check(list) || PyTuple_Check(list));
+
+    struct ast_state *state = get_ast_state();
+    if (state == NULL) {
+        return NULL;
+    }
+
+    Py_ssize_t length = PySequence_Size(list);
+    if (length < 0) {
+        return NULL;
+    }
+    else if (length == 0) {
+        return PyObject_Repr(list);
+    }
+
+    _PyUnicodeWriter writer;
+    _PyUnicodeWriter_Init(&writer);
+    writer.overallocate = 1;
+    PyObject *items[2] = {NULL, NULL};
+
+    items[0] = PySequence_GetItem(list, 0);
+    if (!items[0]) {
+        goto error;
+    }
+    if (length > 1) {
+        items[1] = PySequence_GetItem(list, length - 1);
+        if (!items[1]) {
+            goto error;
+        }
+    }
+
+    bool is_list = PyList_Check(list);
+    if (_PyUnicodeWriter_WriteChar(&writer, is_list ? '[' : '(') < 0) {
+        goto error;
+    }
+
+    for (Py_ssize_t i = 0; i < Py_MIN(length, 2); i++) {
+        PyObject *item = items[i];
+        PyObject *item_repr;
+
+        if (PyType_IsSubtype(Py_TYPE(item), (PyTypeObject *)state->AST_type)) {
+            item_repr = ast_repr_max_depth((AST_object*)item, depth - 1);
+        } else {
+            item_repr = PyObject_Repr(item);
+        }
+        if (!item_repr) {
+            goto error;
+        }
+        if (i > 0) {
+            if (_PyUnicodeWriter_WriteASCIIString(&writer, ", ", 2) < 0) {
+                goto error;
+            }
+        }
+        if (_PyUnicodeWriter_WriteStr(&writer, item_repr) < 0) {
+            Py_DECREF(item_repr);
+            goto error;
+        }
+        if (i == 0 && length > 2) {
+            if (_PyUnicodeWriter_WriteASCIIString(&writer, ", ...", 5) < 0) {
+                Py_DECREF(item_repr);
+                goto error;
+            }
+        }
+        Py_DECREF(item_repr);
+    }
+
+    if (_PyUnicodeWriter_WriteChar(&writer, is_list ? ']' : ')') < 0) {
+        goto error;
+    }
+
+    Py_XDECREF(items[0]);
+    Py_XDECREF(items[1]);
+    return _PyUnicodeWriter_Finish(&writer);
+
+error:
+    Py_XDECREF(items[0]);
+    Py_XDECREF(items[1]);
+    _PyUnicodeWriter_Dealloc(&writer);
+    return NULL;
+}
+
+static PyObject *
+ast_repr_max_depth(AST_object *self, int depth)
+{
+    struct ast_state *state = get_ast_state();
+    if (state == NULL) {
+        return NULL;
+    }
+
+    if (depth <= 0) {
+        return PyUnicode_FromFormat("%s(...)", Py_TYPE(self)->tp_name);
+    }
+
+    int status = Py_ReprEnter((PyObject *)self);
+    if (status != 0) {
+        if (status < 0) {
+            return NULL;
+        }
+        return PyUnicode_FromFormat("%s(...)", Py_TYPE(self)->tp_name);
+    }
+
+    PyObject *fields;
+    if (PyObject_GetOptionalAttr((PyObject *)Py_TYPE(self), state->_fields, &fields) < 0) {
+        Py_ReprLeave((PyObject *)self);
+        return NULL;
+    }
+
+    Py_ssize_t numfields = PySequence_Size(fields);
+    if (numfields < 0) {
+        Py_ReprLeave((PyObject *)self);
+        Py_DECREF(fields);
+        return NULL;
+    }
+
+    if (numfields == 0) {
+        Py_ReprLeave((PyObject *)self);
+        Py_DECREF(fields);
+        return PyUnicode_FromFormat("%s()", Py_TYPE(self)->tp_name);
+    }
+
+    const char* tp_name = Py_TYPE(self)->tp_name;
+    _PyUnicodeWriter writer;
+    _PyUnicodeWriter_Init(&writer);
+    writer.overallocate = 1;
+
+    if (_PyUnicodeWriter_WriteASCIIString(&writer, tp_name, strlen(tp_name)) < 0) {
+        goto error;
+    }
+    if (_PyUnicodeWriter_WriteChar(&writer, '(') < 0) {
+        goto error;
+    }
+
+    for (Py_ssize_t i = 0; i < numfields; i++) {
+        PyObject *name = PySequence_GetItem(fields, i);
+        if (!name) {
+            goto error;
+        }
+
+        PyObject *value = PyObject_GetAttr((PyObject *)self, name);
+        if (!value) {
+            Py_DECREF(name);
+            goto error;
+        }
+
+        PyObject *value_repr;
+        if (PyList_Check(value) || PyTuple_Check(value)) {
+            value_repr = ast_repr_list(value, depth);
+        }
+        else if (PyType_IsSubtype(Py_TYPE(value), (PyTypeObject *)state->AST_type)) {
+            value_repr = ast_repr_max_depth((AST_object*)value, depth - 1);
+        }
+        else {
+            value_repr = PyObject_Repr(value);
+        }
+
+        Py_DECREF(value);
+
+        if (!value_repr) {
+            Py_DECREF(name);
+            Py_DECREF(value);
+            goto error;
+        }
+
+        if (i > 0) {
+            if (_PyUnicodeWriter_WriteASCIIString(&writer, ", ", 2) < 0) {
+                Py_DECREF(name);
+                Py_DECREF(value_repr);
+                goto error;
+            }
+        }
+        if (_PyUnicodeWriter_WriteStr(&writer, name) < 0) {
+            Py_DECREF(name);
+            Py_DECREF(value_repr);
+            goto error;
+        }
+
+        Py_DECREF(name);
+
+        if (_PyUnicodeWriter_WriteChar(&writer, '=') < 0) {
+            Py_DECREF(value_repr);
+            goto error;
+        }
+        if (_PyUnicodeWriter_WriteStr(&writer, value_repr) < 0) {
+            Py_DECREF(value_repr);
+            goto error;
+        }
+
+        Py_DECREF(value_repr);
+    }
+
+    if (_PyUnicodeWriter_WriteChar(&writer, ')') < 0) {
+        goto error;
+    }
+    Py_ReprLeave((PyObject *)self);
+    Py_DECREF(fields);
+    return _PyUnicodeWriter_Finish(&writer);
+
+error:
+    Py_ReprLeave((PyObject *)self);
+    Py_DECREF(fields);
+    _PyUnicodeWriter_Dealloc(&writer);
+    return NULL;
+}
+
+static PyObject *
+ast_repr(AST_object *self)
+{
+    return ast_repr_max_depth(self, 3);
+}
+
 static PyType_Slot AST_type_slots[] = {
     {Py_tp_dealloc, ast_dealloc},
+    {Py_tp_repr, ast_repr},
     {Py_tp_getattro, PyObject_GenericGetAttr},
     {Py_tp_setattro, PyObject_GenericSetAttr},
     {Py_tp_traverse, ast_traverse},