Color the full diff that pytest shows as a diff (#11530)

Related to #11520

Author: BenjaminSchubert

AUTHORS

@@ -56,6 +56,7 @@ Barney Gale
 Ben Gartner
 Ben Webb
 Benjamin Peterson
+Benjamin Schubert
 Bernard Pratz
 Bo Wu
 Bob Ippolito

changelog/11520.improvement.rst

@@ -0,0 +1 @@
+Improved very verbose diff output to color it as a diff instead of only red.

src/_pytest/_io/terminalwriter.py

@@ -3,6 +3,7 @@
 import shutil
 import sys
 from typing import final
+from typing import Literal
 from typing import Optional
 from typing import Sequence
 from typing import TextIO
@@ -193,15 +194,21 @@ def _write_source(self, lines: Sequence[str], indents: Sequence[str] = ()) -> No
         for indent, new_line in zip(indents, new_lines):
             self.line(indent + new_line)
 
-    def _highlight(self, source: str) -> str:
-        """Highlight the given source code if we have markup support."""
+    def _highlight(
+        self, source: str, lexer: Literal["diff", "python"] = "python"
+    ) -> str:
+        """Highlight the given source if we have markup support."""
         from _pytest.config.exceptions import UsageError
 
         if not self.hasmarkup or not self.code_highlight:
             return source
         try:
             from pygments.formatters.terminal import TerminalFormatter
-            from pygments.lexers.python import PythonLexer
+
+            if lexer == "python":
+                from pygments.lexers.python import PythonLexer as Lexer
+            elif lexer == "diff":
+                from pygments.lexers.diff import DiffLexer as Lexer
             from pygments import highlight
             import pygments.util
         except ImportError:
@@ -210,7 +217,7 @@ def _highlight(self, source: str) -> str:
             try:
                 highlighted: str = highlight(
                     source,
-                    PythonLexer(),
+                    Lexer(),
                     TerminalFormatter(
                         bg=os.getenv("PYTEST_THEME_MODE", "dark"),
                         style=os.getenv("PYTEST_THEME"),

src/_pytest/assertion/util.py

@@ -7,8 +7,10 @@
 from typing import Callable
 from typing import Iterable
 from typing import List
+from typing import Literal
 from typing import Mapping
 from typing import Optional
+from typing import Protocol
 from typing import Sequence
 from unicodedata import normalize
 
@@ -33,6 +35,11 @@
 _config: Optional[Config] = None
 
 
+class _HighlightFunc(Protocol):
+    def __call__(self, source: str, lexer: Literal["diff", "python"] = "python") -> str:
+        """Apply highlighting to the given source."""
+
+
 def format_explanation(explanation: str) -> str:
     r"""Format an explanation.
 
@@ -189,7 +196,8 @@ def assertrepr_compare(
     explanation = None
     try:
         if op == "==":
-            explanation = _compare_eq_any(left, right, verbose)
+            writer = config.get_terminal_writer()
+            explanation = _compare_eq_any(left, right, writer._highlight, verbose)
         elif op == "not in":
             if istext(left) and istext(right):
                 explanation = _notin_text(left, right, verbose)
@@ -225,7 +233,9 @@ def assertrepr_compare(
     return [summary] + explanation
 
 
-def _compare_eq_any(left: Any, right: Any, verbose: int = 0) -> List[str]:
+def _compare_eq_any(
+    left: Any, right: Any, highlighter: _HighlightFunc, verbose: int = 0
+) -> List[str]:
     explanation = []
     if istext(left) and istext(right):
         explanation = _diff_text(left, right, verbose)
@@ -245,7 +255,7 @@ def _compare_eq_any(left: Any, right: Any, verbose: int = 0) -> List[str]:
             # field values, not the type or field names. But this branch
             # intentionally only handles the same-type case, which was often
             # used in older code bases before dataclasses/attrs were available.
-            explanation = _compare_eq_cls(left, right, verbose)
+            explanation = _compare_eq_cls(left, right, highlighter, verbose)
         elif issequence(left) and issequence(right):
             explanation = _compare_eq_sequence(left, right, verbose)
         elif isset(left) and isset(right):
@@ -254,7 +264,7 @@ def _compare_eq_any(left: Any, right: Any, verbose: int = 0) -> List[str]:
             explanation = _compare_eq_dict(left, right, verbose)
 
         if isiterable(left) and isiterable(right):
-            expl = _compare_eq_iterable(left, right, verbose)
+            expl = _compare_eq_iterable(left, right, highlighter, verbose)
             explanation.extend(expl)
 
     return explanation
@@ -321,7 +331,10 @@ def _surrounding_parens_on_own_lines(lines: List[str]) -> None:
 
 
 def _compare_eq_iterable(
-    left: Iterable[Any], right: Iterable[Any], verbose: int = 0
+    left: Iterable[Any],
+    right: Iterable[Any],
+    highligher: _HighlightFunc,
+    verbose: int = 0,
 ) -> List[str]:
     if verbose <= 0 and not running_on_ci():
         return ["Use -v to get more diff"]
@@ -346,7 +359,13 @@ def _compare_eq_iterable(
     # "right" is the expected base against which we compare "left",
     # see https://github.com/pytest-dev/pytest/issues/3333
     explanation.extend(
-        line.rstrip() for line in difflib.ndiff(right_formatting, left_formatting)
+        highligher(
+            "\n".join(
+                line.rstrip()
+                for line in difflib.ndiff(right_formatting, left_formatting)
+            ),
+            lexer="diff",
+        ).splitlines()
     )
     return explanation
 
@@ -496,7 +515,9 @@ def _compare_eq_dict(
     return explanation
 
 
-def _compare_eq_cls(left: Any, right: Any, verbose: int) -> List[str]:
+def _compare_eq_cls(
+    left: Any, right: Any, highlighter: _HighlightFunc, verbose: int
+) -> List[str]:
     if not has_default_eq(left):
         return []
     if isdatacls(left):
@@ -542,7 +563,9 @@ def _compare_eq_cls(left: Any, right: Any, verbose: int) -> List[str]:
             ]
             explanation += [
                 indent + line
-                for line in _compare_eq_any(field_left, field_right, verbose)
+                for line in _compare_eq_any(
+                    field_left, field_right, highlighter, verbose
+                )
             ]
     return explanation
 

testing/conftest.py

@@ -160,6 +160,9 @@ class ColorMapping:
             "red": "\x1b[31m",
             "green": "\x1b[32m",
             "yellow": "\x1b[33m",
+            "light-gray": "\x1b[90m",
+            "light-red": "\x1b[91m",
+            "light-green": "\x1b[92m",
             "bold": "\x1b[1m",
             "reset": "\x1b[0m",
             "kw": "\x1b[94m",
@@ -171,6 +174,7 @@ class ColorMapping:
             "endline": "\x1b[90m\x1b[39;49;00m",
         }
         RE_COLORS = {k: re.escape(v) for k, v in COLORS.items()}
+        NO_COLORS = {k: "" for k in COLORS.keys()}
 
         @classmethod
         def format(cls, lines: List[str]) -> List[str]:
@@ -187,6 +191,11 @@ def format_for_rematch(cls, lines: List[str]) -> List[str]:
             """Replace color names for use with LineMatcher.re_match_lines"""
             return [line.format(**cls.RE_COLORS) for line in lines]
 
+        @classmethod
+        def strip_colors(cls, lines: List[str]) -> List[str]:
+            """Entirely remove every color code"""
+            return [line.format(**cls.NO_COLORS) for line in lines]
+
     return ColorMapping
 
 

testing/test_assertion.py

@@ -18,12 +18,19 @@
 
 
 def mock_config(verbose=0):
+    class TerminalWriter:
+        def _highlight(self, source, lexer):
+            return source
+
     class Config:
         def getoption(self, name):
             if name == "verbose":
                 return verbose
             raise KeyError("Not mocked out: %s" % name)
 
+        def get_terminal_writer(self):
+            return TerminalWriter()
+
     return Config()
 
 
@@ -1784,3 +1791,48 @@ def test_reprcompare_verbose_long() -> None:
         "{'v0': 0, 'v1': 1, 'v2': 12, 'v3': 3, 'v4': 4, 'v5': 5, "
         "'v6': 6, 'v7': 7, 'v8': 8, 'v9': 9, 'v10': 10}"
     )
+
+
+@pytest.mark.parametrize("enable_colors", [True, False])
+@pytest.mark.parametrize(
+    ("test_code", "expected_lines"),
+    (
+        (
+            """
+            def test():
+                assert [0, 1] == [0, 2]
+            """,
+            [
+                "{bold}{red}E         {light-red}- [0, 2]{hl-reset}{endline}{reset}",
+                "{bold}{red}E         {light-green}+ [0, 1]{hl-reset}{endline}{reset}",
+            ],
+        ),
+        (
+            """
+            def test():
+                assert {f"number-is-{i}": i for i in range(1, 6)} == {
+                    f"number-is-{i}": i for i in range(5)
+                }
+            """,
+            [
+                "{bold}{red}E         {light-gray} {hl-reset} {{{endline}{reset}",
+                "{bold}{red}E         {light-gray} {hl-reset}  'number-is-1': 1,{endline}{reset}",
+                "{bold}{red}E         {light-green}+  'number-is-5': 5,{hl-reset}{endline}{reset}",
+            ],
+        ),
+    ),
+)
+def test_comparisons_handle_colors(
+    pytester: Pytester, color_mapping, enable_colors, test_code, expected_lines
+) -> None:
+    p = pytester.makepyfile(test_code)
+    result = pytester.runpytest(
+        f"--color={'yes' if enable_colors else 'no'}", "-vv", str(p)
+    )
+    formatter = (
+        color_mapping.format_for_fnmatch
+        if enable_colors
+        else color_mapping.strip_colors
+    )
+
+    result.stdout.fnmatch_lines(formatter(expected_lines), consecutive=False)