Add colour to pprint output

Author: hugovk

Doc/library/pprint.rst

@@ -31,7 +31,8 @@ Functions
 ---------
 
 .. function:: pp(object, stream=None, indent=1, width=80, depth=None, *, \
-                     compact=False, sort_dicts=False, underscore_numbers=False)
+                     color=True, compact=False, sort_dicts=False, \
+                     underscore_numbers=False)
 
    Prints the formatted representation of *object*, followed by a newline.
    This function may be used in the interactive interpreter
@@ -63,6 +64,12 @@ Functions
       on the depth of the objects being formatted.
    :type depth: int | None
 
+   :param bool color:
+      If ``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.
+
    :param bool compact:
       Control the way long :term:`sequences <sequence>` are formatted.
       If ``False`` (the default),
@@ -93,14 +100,21 @@ Functions
 
    .. versionadded:: 3.8
 
+   .. versionchanged:: next
+      Added the *color* parameter.
+
 
 .. function:: pprint(object, stream=None, indent=1, width=80, depth=None, *, \
-                     compact=False, sort_dicts=True, underscore_numbers=False)
+                     color=True, compact=False, sort_dicts=True, \
+                     underscore_numbers=False)
 
    Alias for :func:`~pprint.pp` with *sort_dicts* set to ``True`` by default,
    which would automatically sort the dictionaries' keys,
    you might want to use :func:`~pprint.pp` instead where it is ``False`` by default.
 
+   .. versionchanged:: next
+      Added the *color* parameter.
+
 
 .. function:: pformat(object, indent=1, width=80, depth=None, *, \
                       compact=False, sort_dicts=True, underscore_numbers=False)
@@ -144,13 +158,14 @@ Functions
 
 .. _prettyprinter-objects:
 
-PrettyPrinter Objects
+PrettyPrinter objects
 ---------------------
 
 .. index:: single: ...; placeholder
 
 .. class:: PrettyPrinter(indent=1, width=80, depth=None, stream=None, *, \
-                         compact=False, sort_dicts=True, underscore_numbers=False)
+                         color=True, compact=False, sort_dicts=True, \
+                         underscore_numbers=False)
 
    Construct a :class:`PrettyPrinter` instance.
 
@@ -193,6 +208,9 @@ PrettyPrinter Objects
    .. versionchanged:: 3.11
       No longer attempts to write to :data:`!sys.stdout` if it is ``None``.
 
+   .. versionchanged:: next
+      Added the *color* parameter.
+
 
 :class:`PrettyPrinter` instances have the following methods:
 

Doc/whatsnew/3.15.rst

@@ -864,6 +864,16 @@ pickle
   (Contributed by Zackery Spytz and Serhiy Storchaka in :gh:`77188`.)
 
 
+pprint
+------
+
+* Add *color* parameter to :func:`~pprint.pp` and :func:`~pprint.pprint`.
+  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 Hugo van Kemenade in :gh:`123456`.)
+
+
 re
 --
 

Lib/pprint.py

@@ -39,17 +39,36 @@
 import types as _types
 from io import StringIO as _StringIO
 
+lazy import _colorize
+lazy from _pyrepl.utils import disp_str, gen_colors
+
 __all__ = ["pprint","pformat","isreadable","isrecursive","saferepr",
            "PrettyPrinter", "pp"]
 
 
-def pprint(object, stream=None, indent=1, width=80, depth=None, *,
-           compact=False, sort_dicts=True, underscore_numbers=False):
+def pprint(
+    object,
+    stream=None,
+    indent=1,
+    width=80,
+    depth=None,
+    *,
+    color=True,
+    compact=False,
+    sort_dicts=True,
+    underscore_numbers=False,
+):
     """Pretty-print a Python object to a stream [default is sys.stdout]."""
     printer = PrettyPrinter(
-        stream=stream, indent=indent, width=width, depth=depth,
-        compact=compact, sort_dicts=sort_dicts,
-        underscore_numbers=underscore_numbers)
+        stream=stream,
+        indent=indent,
+        width=width,
+        depth=depth,
+        color=color,
+        compact=compact,
+        sort_dicts=sort_dicts,
+        underscore_numbers=underscore_numbers,
+    )
     printer.pprint(object)
 
 
@@ -109,9 +128,26 @@ def _safe_tuple(t):
     return _safe_key(t[0]), _safe_key(t[1])
 
 
+def _colorize_output(text):
+    """Apply syntax highlighting."""
+    colors = list(gen_colors(text))
+    chars, _ = disp_str(text, colors=colors, force_color=True)
+    return "".join(chars)
+
+
 class PrettyPrinter:
-    def __init__(self, indent=1, width=80, depth=None, stream=None, *,
-                 compact=False, sort_dicts=True, underscore_numbers=False):
+    def __init__(
+        self,
+        indent=1,
+        width=80,
+        depth=None,
+        stream=None,
+        *,
+        color=True,
+        compact=False,
+        sort_dicts=True,
+        underscore_numbers=False,
+    ):
         """Handle pretty printing operations onto a stream using a set of
         configured parameters.
 
@@ -128,6 +164,11 @@ def __init__(self, indent=1, width=80, depth=None, stream=None, *,
             The desired output stream.  If omitted (or false), the standard
             output stream available at construction will be used.
 
+        color
+            If true (the default), syntax highlighting is enabled for pprint
+            when the stream and environment variables permit.
+            If false, colored output is always disabled.
+
         compact
             If true, several items will be combined in one line.
 
@@ -156,10 +197,16 @@ def __init__(self, indent=1, width=80, depth=None, stream=None, *,
         self._compact = bool(compact)
         self._sort_dicts = sort_dicts
         self._underscore_numbers = underscore_numbers
+        self._color = color
 
     def pprint(self, object):
         if self._stream is not None:
-            self._format(object, self._stream, 0, 0, {}, 0)
+            if self._color and _colorize.can_colorize(file=self._stream):
+                sio = _StringIO()
+                self._format(object, sio, 0, 0, {}, 0)
+                self._stream.write(_colorize_output(sio.getvalue()))
+            else:
+                self._format(object, self._stream, 0, 0, {}, 0)
             self._stream.write("\n")
 
     def pformat(self, object):

Lib/test/test_pickle.py

@@ -779,6 +779,7 @@ def invoke_pickle(self, *flags):
             pickle._main(args=[*flags, self.filename])
         return self.text_normalize(output.getvalue())
 
+    @support.force_not_colorized
     def test_invocation(self):
         # test 'python -m pickle pickle_file'
         data = {

Lib/test/test_pprint.py

@@ -10,6 +10,7 @@
 import re
 import types
 import unittest
+import unittest.mock
 from collections.abc import ItemsView, KeysView, Mapping, MappingView, ValuesView
 
 from test.support import cpython_only
@@ -165,6 +166,53 @@ def test_init(self):
         self.assertRaises(ValueError, pprint.PrettyPrinter, depth=-1)
         self.assertRaises(ValueError, pprint.PrettyPrinter, width=0)
 
+    def test_color_pprint(self):
+        """Test pprint color parameter."""
+        obj = {"key": "value"}
+        stream = io.StringIO()
+
+        # color=False should produce no ANSI codes
+        pprint.pprint(obj, stream=stream, color=False)
+        result = stream.getvalue()
+        self.assertNotIn("\x1b[", result)
+
+        # Explicit color=False should override FORCE_COLOR
+        stream = io.StringIO()
+        with unittest.mock.patch.dict(
+            "os.environ", {"FORCE_COLOR": "1", "NO_COLOR": ""}
+        ):
+            pprint.pprint(obj, stream=stream, color=False)
+            result = stream.getvalue()
+            self.assertNotIn("\x1b[", result)
+
+    def test_color_prettyprinter(self):
+        """Test PrettyPrinter color parameter."""
+        obj = {"key": "value"}
+
+        # color=False should produce no ANSI codes in pprint
+        stream = io.StringIO()
+        pp = pprint.PrettyPrinter(stream=stream, color=False)
+        pp.pprint(obj)
+        self.assertNotIn("\x1b[", stream.getvalue())
+
+        # color=True with FORCE_COLOR should produce ANSI codes in pprint
+        with unittest.mock.patch.dict(
+            "os.environ", {"FORCE_COLOR": "1", "NO_COLOR": ""}
+        ):
+            stream = io.StringIO()
+            pp = pprint.PrettyPrinter(stream=stream, color=True)
+            pp.pprint(obj)
+            self.assertIn("\x1b[", stream.getvalue())
+
+        # Explicit color=False should override FORCE_COLOR
+        with unittest.mock.patch.dict(
+            "os.environ", {"FORCE_COLOR": "1", "NO_COLOR": ""}
+        ):
+            stream = io.StringIO()
+            pp = pprint.PrettyPrinter(stream=stream, color=False)
+            pp.pprint(obj)
+            self.assertNotIn("\x1b[", stream.getvalue())
+
     def test_basic(self):
         # Verify .isrecursive() and .isreadable() w/o recursion
         pp = pprint.PrettyPrinter()