gh-145217: Add colour to pprint output

[hugovk]

Example:

import pprint

data = ["spam", "eggs", ("lumberjack", "knights"), "ni", 42, None]
data.insert(0, data)

pprint.pp(data)
pprint.pprint(data)
pprint.PrettyPrinter().pprint(data)

• Issue: Add colour to pprint output #145217

---

📚 Documentation preview 📚: https://cpython-previews--145218.org.readthedocs.build/

[sunmy2019]

Shall we add a bool(color) conversion just as a few lines above?

        self._compact = bool(compact)
        self._expand = bool(expand)

[hugovk]

Sure, can do, but not sure if it's necessary?

[sunmy2019]

Do we have explicit test for color=True and can_colorize=False? It's an end to end test.

something like

        obj = {"key": "value"}
        stream = io.StringIO()

        # color=True should not produce no ANSI codes for streams
        # that do not support color
        pprint.pprint(obj, stream=stream, color=True)
        result = stream.getvalue()
        self.assertNotIn("\x1b[", result)

[sunmy2019]

If user's __repr__ contains color, it would break the color.

Code

import os
import pprint

os.environ["FORCE_COLOR"] = "1"
os.environ["NO_COLOR"] = ""


class Colorful:
    """something with a colorful repr"""

    def __repr__(self):
        return "\x1b[31mred\x1b[0m"

pprint.pprint({"a": Colorful(), "b": 42, "c": "hello"}, color=True)

What should we do for this case?

1. Document and keep as is?
2. Pre check escape sequence before colorize?
3. Other techniques for this?

[hugovk]

2. Pre check escape sequence before colorize?

I think if there's any colour escapes, we bail and print it as-is without extra colour from us. So we get something like this:

And I expect colour is rare in users' __repr__s? If so, this won't happen often, and when it does, we don't override any of their own colour or need more complex handling. And they don't end up with just a green "a" like in your example.

[read-the-docs-community]

Documentation build overview

📚 cpython-previews | 🛠️ Build #32469897 | 📁 Comparing 1d05505 against main (4b33308)

  🔍 Preview build  

6 files changed · ± 6 modified

± Modified

• contents.html
• download.html
• library/datatypes.html
• library/pprint.html
• whatsnew/3.15.html
• whatsnew/changelog.html

[sunmy2019]

I think if there's any colour escapes, we bail and print it as-is without extra colour from us.

I agree with this.

But thinking further, our current highlighting logic relies on tokenize, yet pprint output has never made any guarantee that it must be valid Python source. __repr__ can return anything — it doesn't have to be tokenizer-friendly.

So, instead of special-casing \x1b[, how about

Probe the text with tokenize first,
-- if it can process the whole thing, highlight it;
-- if it chokes halfway through for any reason (ANSI escapes, illegal control chars, etc.), just print the raw text with no added colors.

No blacklist needed, no half-highlighted output.

What do you think?

 def _colorize_output(text):
     """Apply syntax highlighting."""
-    if "\x1b[" in text:
-        # If the text already contains ANSI escape sequences
-        # (for example, from a custom __repr__),
-        # return as-is to avoid breaking their color.
+    # Probe: if the text cannot be fully tokenized, fall back to plain text
+    # to avoid the "half-highlighted" problem.
+    try:
+        list(tokenize.generate_tokens(_StringIO(text).readline))
+    except (SyntaxError, tokenize.TokenError):
         return text
     colors = list(gen_colors(text))
     chars, _ = disp_str(text, colors=colors, force_color=True, escape=False)

[hugovk]

Could do, although that'd mean we're potentially doing tokenize.generate_tokens() both here and in gen_colors(). Maybe that doesn't matter.

[sunmy2019]

Maybe that doesn't matter.

If that matters, we could modify the latter to raise errors (rather than catch them)

[hugovk]

I think I'll leave it like this, unless others also want to suggest something. We can always fix it later if needed. Thanks for the review!