gh-148981: Add color parameter to ast.dump (#148982)

And turn on color for the `ast` module CLI.

Author: StanFromIreland

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=False, 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,10 @@ 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 returned string is syntax highlighted using
+   ANSI escape sequences.
+   If ``False`` (the default), 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 +2531,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:
 
@@ -2584,6 +2591,10 @@ Command-line usage
 
 .. versionadded:: 3.9
 
+.. versionchanged:: next
+   The output is now syntax highlighted by default. This can be
+   :ref:`controlled using environment variables <using-on-controlling-color>`.
+
 The :mod:`!ast` module can be executed as a script from the command line.
 It is as simple as:
 

Doc/whatsnew/3.15.rst

@@ -704,6 +704,20 @@ array
   (Contributed by Sergey B Kirpichev in :gh:`146238`.)
 
 
+ast
+---
+
+* Add *color* parameter to :func:`~ast.dump`.
+  If ``True``, the returned string is syntax highlighted using ANSI escape
+  sequences.
+  If ``False`` (the default), colored output is always disabled.
+  (Contributed by Stan Ulbrych in :gh:`148981`.)
+
+* The :ref:`command-line <ast-cli>` output is now syntax highlighted by default.
+  This can be :ref:`controlled using environment variables <using-on-controlling-color>`.
+  (Contributed by Stan Ulbrych in :gh:`148981`.)
+
+
 base64
 ------
 

Lib/_colorize.py

@@ -189,6 +189,17 @@ class Argparse(ThemeSection):
     message: str = ANSIColors.MAGENTA
 
 
+@dataclass(frozen=True, kw_only=True)
+class Ast(ThemeSection):
+    node: str = ANSIColors.CYAN
+    field: str = ANSIColors.BLUE
+    attribute: str = ANSIColors.GREY
+    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 +416,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 +430,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 +447,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 +468,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,7 @@
 :license: Python License.
 """
 from _ast import *
+lazy from _colorize import can_colorize, get_theme
 
 
 def parse(source, filename='<unknown>', mode='exec', *,
@@ -117,21 +118,32 @@ def _convert_literal(node):
 def dump(
     node, annotate_fields=True, include_attributes=False,
     *,
-    indent=None, show_empty=False,
+    color=False, 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),
-    the returned string will show the names and the values for fields.
-    If annotate_fields is false, the result string will be more compact by
-    omitting unambiguous field names.  Attributes such as line
-    numbers and column offsets are not dumped by default.  If this is wanted,
-    include_attributes can be set to true.  If indent is a non-negative
-    integer or string, then the tree will be pretty-printed with that indent
-    level. None (the default) selects the single line representation.
+    debugging purposes.
+
+    If annotate_fields is true (by default), the returned string will show the
+    names and the values for fields. If annotate_fields is false, the result
+    string will be more compact by omitting unambiguous field names.
+
+    Attributes such as line 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 returned string is syntax highlighted using ANSI
+    escape sequences. If color is false (the default), 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. If indent is None (the default),
+    the tree is dumped on a single line.
+
     If show_empty is False, then empty lists and fields that are None
     will be omitted from the output for better readability.
     """
+    t = get_theme(force_color=color, force_no_color=not color).ast
+
     def _format(node, level=0):
         if indent is not None:
             level += 1
@@ -166,15 +178,17 @@ def _format(node, level=0):
                         field_type = cls._field_types.get(name, object)
                         if field_type is expr_context:
                             if not keywords:
-                                args_buffer.append(repr(value))
+                                args_buffer.append(
+                                    f'{t.node}{type(value).__name__}'
+                                    f'{t.reset}()')
                             continue
                     if not keywords:
                         args.extend(args_buffer)
                         args_buffer = []
                 value, simple = _format(value, level)
                 allsimple = allsimple and simple
                 if keywords:
-                    args.append('%s=%s' % (name, value))
+                    args.append(f'{t.field}{name}{t.reset}={value}')
                 else:
                     args.append(value)
             if include_attributes and node._attributes:
@@ -187,14 +201,21 @@ def _format(node, level=0):
                         continue
                     value, simple = _format(value, level)
                     allsimple = allsimple and simple
-                    args.append('%s=%s' % (name, value))
+                    args.append(f'{t.attribute}{name}{t.reset}={value}')
+            cls_name = f'{t.node}{cls.__name__}{t.reset}'
             if allsimple and len(args) <= 3:
-                return '%s(%s)' % (node.__class__.__name__, ', '.join(args)), not args
-            return '%s(%s%s)' % (node.__class__.__name__, prefix, sep.join(args)), False
+                return f'{cls_name}({", ".join(args)})', not args
+            return f'{cls_name}({prefix}{sep.join(args)})', False
         elif isinstance(node, list):
             if not node:
                 return '[]', True
             return '[%s%s]' % (prefix, sep.join(_format(x, level)[0] for x in node)), False
+        if isinstance(node, bool) or node is None or node is Ellipsis:
+            return f'{t.keyword}{node!r}{t.reset}', True
+        if isinstance(node, (int, float, complex)):
+            return f'{t.number}{node!r}{t.reset}', True
+        if isinstance(node, (str, bytes)):
+            return f'{t.string}{node!r}{t.reset}', True
         return repr(node), True
 
     if not isinstance(node, AST):
@@ -642,7 +663,7 @@ def main(args=None):
     import argparse
     import sys
 
-    parser = argparse.ArgumentParser(color=True)
+    parser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter)
     parser.add_argument('infile', nargs='?', default='-',
                         help='the file to parse; defaults to stdin')
     parser.add_argument('-m', '--mode', default='exec',
@@ -661,7 +682,7 @@ def main(args=None):
                              '(for example, 3.10)')
     parser.add_argument('-O', '--optimize',
                         type=int, default=-1, metavar='LEVEL',
-                        help='optimization level for parser (default -1)')
+                        help='optimization level for parser')
     parser.add_argument('--show-empty', default=False, action='store_true',
                         help='show empty lists and fields in dump output')
     args = parser.parse_args(args)
@@ -688,6 +709,7 @@ def main(args=None):
     tree = parse(source, name, args.mode, type_comments=args.no_type_comments,
                  feature_version=feature_version, optimize=args.optimize)
     print(dump(tree, include_attributes=args.include_attributes,
+               color=can_colorize(file=sys.stdout),
                indent=args.indent, show_empty=args.show_empty))
 
 if __name__ == '__main__':

Lib/test/test_ast/test_ast.py

@@ -1705,6 +1705,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")
+        self.assertNotIn("\x1b[", ast.dump(node))
+        self.assertNotIn("\x1b[", ast.dump(node, color=False))
+        self.assertIn("\x1b[", ast.dump(node, color=True))
+
+        node = ast.Constant(value="\x1b[31m")
+        self.assertEqual(ast.dump(node), "Constant(value='\\x1b[31m')")
+        self.assertIn("'\\x1b[31m'", 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 +3425,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()

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`.