gh-148981: Add color to ast.dump

Doc/library/ast.rst

@@ -2480,7 +2480,7 @@ and classes for traversing abstract syntax trees:
       node = YourTransformer().visit(node)
 
 
-.. function:: dump(node, annotate_fields=True, include_attributes=False, *, indent=None, show_empty=False)
+.. function:: dump(node, annotate_fields=True, include_attributes=False, *, color=True, indent=None, show_empty=False)
 
    Return a formatted dump of the tree in *node*.  This is mainly useful for
    debugging purposes.  If *annotate_fields* is true (by default),
@@ -2490,6 +2490,11 @@ and classes for traversing abstract syntax trees:
    numbers and column offsets are not dumped by default.  If this is wanted,
    *include_attributes* can be set to true.
 
+   If *color* is ``True`` (the default), output will be syntax highlighted using
+   ANSI escape sequences, if the *stream* and :ref:`environment variables
+   <using-on-controlling-color>` permit.
+   If ``False``, colored output is always disabled.
+
    If *indent* is a non-negative integer or string, then the tree will be
    pretty-printed with that indent level.  An indent level
    of 0, negative, or ``""`` will only insert newlines.  ``None`` (the default)
@@ -2527,6 +2532,9 @@ and classes for traversing abstract syntax trees:
    .. versionchanged:: 3.15
       Omit optional ``Load()`` values by default.
 
+   .. versionchanged:: next
+      Added the *color* parameter.
+
 
 .. _ast-compiler-flags:
 

Doc/whatsnew/3.15.rst

@@ -692,6 +692,16 @@ array
   (Contributed by Sergey B Kirpichev in :gh:`146238`.)
 
 
+ast
+---
+
+* Add *color* parameter to :func:`~ast.dump`.
+  If ``True`` (the default), output is highlighted in color, when the stream
+  and :ref:`environment variables <using-on-controlling-color>` permit.
+  If ``False``, colored output is always disabled.
+  (Contributed by Stan Ulbrych in :gh:`148981`.)
+
+
 base64
 ------
 

Lib/_colorize.py

@@ -189,6 +189,16 @@ class Argparse(ThemeSection):
     message: str = ANSIColors.MAGENTA
 
 
+@dataclass(frozen=True, kw_only=True)
+class Ast(ThemeSection):
+    node: str = ANSIColors.CYAN
+    field: str = ANSIColors.BLUE
+    string: str = ANSIColors.GREEN
+    number: str = ANSIColors.YELLOW
+    keyword: str = ANSIColors.BOLD_BLUE
+    reset: str = ANSIColors.RESET
+
+
 @dataclass(frozen=True, kw_only=True)
 class Difflib(ThemeSection):
     """A 'git diff'-like theme for `difflib.unified_diff`."""
@@ -405,6 +415,7 @@ class Theme:
     below.
     """
     argparse: Argparse = field(default_factory=Argparse)
+    ast: Ast = field(default_factory=Ast)
     difflib: Difflib = field(default_factory=Difflib)
     fancycompleter: FancyCompleter = field(default_factory=FancyCompleter)
     http_server: HttpServer = field(default_factory=HttpServer)
@@ -418,6 +429,7 @@ def copy_with(
         self,
         *,
         argparse: Argparse | None = None,
+        ast: Ast | None = None,
         difflib: Difflib | None = None,
         fancycompleter: FancyCompleter | None = None,
         http_server: HttpServer | None = None,
@@ -434,6 +446,7 @@ def copy_with(
         """
         return type(self)(
             argparse=argparse or self.argparse,
+            ast=ast or self.ast,
             difflib=difflib or self.difflib,
             fancycompleter=fancycompleter or self.fancycompleter,
             http_server=http_server or self.http_server,
@@ -454,6 +467,7 @@ def no_colors(cls) -> Self:
         """
         return cls(
             argparse=Argparse.no_colors(),
+            ast=Ast.no_colors(),
             difflib=Difflib.no_colors(),
             fancycompleter=FancyCompleter.no_colors(),
             http_server=HttpServer.no_colors(),

Lib/ast.py

@@ -21,6 +21,9 @@
 :license: Python License.
 """
 from _ast import *
+lazy import re
+lazy import sys
+lazy from _colorize import can_colorize, get_theme
 
 
 def parse(source, filename='<unknown>', mode='exec', *,
@@ -117,7 +120,7 @@ def _convert_literal(node):
 def dump(
     node, annotate_fields=True, include_attributes=False,
     *,
-    indent=None, show_empty=False,
+    color=True, indent=None, show_empty=False,
 ):
     """
     Return a formatted dump of the tree in node.  This is mainly useful for
@@ -131,6 +134,9 @@ def dump(
     level. None (the default) selects the single line representation.
     If show_empty is False, then empty lists and fields that are None
     will be omitted from the output for better readability.
+    If color is true (the default), the result will be syntax highlighted
+    using ANSI escape sequences if the stream and environment variables permit.
+    If color is false, colored output is always disabled.
     """
     def _format(node, level=0):
         if indent is not None:
@@ -201,7 +207,38 @@ def _format(node, level=0):
         raise TypeError('expected AST, got %r' % node.__class__.__name__)
     if indent is not None and not isinstance(indent, str):
         indent = ' ' * indent
-    return _format(node)[0]
+    output = _format(node)[0]
+    if color:
+        if can_colorize(file=sys.stdout):
+            output = _colorize_dump(output, get_theme(tty_file=sys.stdout).ast)
+    return output
+
+
+_color_pattern = None
+
+def _colorize_dump(output, theme):
+    global _color_pattern
+    if _color_pattern is None:
+        _color_pattern = re.compile(r"""
+            (?P<string>[bB]?'(?:\\.|[^'\\])*'|[bB]?"(?:\\.|[^"\\])*") |
+            (?P<keyword>\b(?:None|True|False|Ellipsis)\b)             |
+            (?P<node>[A-Za-z_]\w*)(?=\()                              |
+            (?P<field>[a-z_]\w*)(?==)                                 |
+            (?P<number>-?\d+(?:\.\d+)?(?:[eE][+-]?\d+)?[jJ]?)
+        """, re.VERBOSE)
+
+    color_map = {
+        "node": theme.node,
+        "field": theme.field,
+        "string": theme.string,
+        "number": theme.number,
+        "keyword": theme.keyword,
+    }
+
+    def replace(match):
+        return f"{color_map[match.lastgroup]}{match.group()}{theme.reset}"
+
+    return _color_pattern.sub(replace, output)
 
 
 def copy_location(new_node, old_node):
@@ -337,8 +374,6 @@ def _splitlines_no_ff(source, maxlines=None):
     """
     global _line_pattern
     if _line_pattern is None:
-        # lazily computed to speedup import time of `ast`
-        import re
         _line_pattern = re.compile(r"(.*?(?:\r\n|\n|\r|$))")
 
     lines = []
@@ -640,7 +675,6 @@ def unparse(ast_obj):
 
 def main(args=None):
     import argparse
-    import sys
 
     parser = argparse.ArgumentParser(color=True)
     parser.add_argument('infile', nargs='?', default='-',

Lib/test/test_ast/test_ast.py

@@ -1450,6 +1450,7 @@ def test_replace_non_str_kwarg(self):
             node.__replace__(**{object(): "y"})
 
 
+@support.force_not_colorized_test_class
 class ASTHelpers_Test(unittest.TestCase):
     maxDiff = None
 
@@ -1705,6 +1706,16 @@ def check_text(code, empty, full, **kwargs):
             full="Module(body=[Import(names=[alias(name='_ast', asname='ast')], is_lazy=0), ImportFrom(module='module', names=[alias(name='sub')], level=0, is_lazy=0)], type_ignores=[])",
         )
 
+    def test_dump_with_color(self):
+        node = ast.parse("x = 1")
+
+        with support.force_color(True):
+            self.assertNotIn("\x1b[", ast.dump(node, color=False))
+            self.assertIn("\x1b[", ast.dump(node, color=True))
+
+        with support.force_color(False):
+            self.assertNotIn("\x1b[", ast.dump(node, color=True))
+
     def test_copy_location(self):
         src = ast.parse('1 + 1', mode='eval')
         src.body.right = ast.copy_location(ast.Constant(2), src.body.right)
@@ -3415,6 +3426,7 @@ def test_subinterpreter(self):
         self.assertEqual(res, 0)
 
 
+@support.force_not_colorized_test_class
 class CommandLineTests(unittest.TestCase):
     def setUp(self):
         self.filename = tempfile.mktemp()
@@ -3674,6 +3686,7 @@ def test_show_empty_flag(self):
         self.check_output(source, expect, '--show-empty')
 
 
+@support.force_not_colorized_test_class
 class ASTOptimizationTests(unittest.TestCase):
     def wrap_expr(self, expr):
         return ast.Module(body=[ast.Expr(value=expr)])

Misc/NEWS.d/next/Library/2026-04-25-12-50-46.gh-issue-148981.YMM4Y9.rst

@@ -0,0 +1 @@
+Add *color* parameter to :func:`ast.dump`.