Don't add colour if there's already colour

Author: hugovk

Lib/pprint.py

@@ -133,6 +133,11 @@ def _safe_tuple(t):
 
 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.
+        return text
     colors = list(gen_colors(text))
     chars, _ = disp_str(text, colors=colors, force_color=True, escape=False)
     return "".join(chars)

Lib/test/test_pprint.py

@@ -245,6 +245,26 @@ def test_color_preserves_newlines(self):
         stripped = re.sub(r"\x1b\[[0-9;]*m", "", color)
         self.assertEqual(stripped, plain)
 
+    def test_color_user_repr_with_ansi(self):
+        """If a __repr__ already contains ANSI escapes, don't add ours."""
+
+        class ColorObj:
+            def __repr__(self):
+                return "\x1b[31mred\x1b[0m"
+
+        obj = {"a": ColorObj(), "b": 42, "c": "hello"}
+
+        with unittest.mock.patch.dict(
+            "os.environ", {"FORCE_COLOR": "1", "NO_COLOR": ""}
+        ):
+            stream = io.StringIO()
+            pprint.pprint(obj, stream=stream, color=True)
+            result = stream.getvalue()
+
+        # pprint should not have added any extra color codes
+        expected = "{'a': \x1b[31mred\x1b[0m, 'b': 42, 'c': 'hello'}\n"
+        self.assertEqual(result, expected)
+
     def test_basic(self):
         # Verify .isrecursive() and .isreadable() w/o recursion
         pp = pprint.PrettyPrinter()