Use code highlighting if pygments is installed (#6658)

* Use code highlighting if pygments is installed

* Use colorama constants instead of bare ascii codes

Could not find the exact equivalent of 'hl-reset' code using colorama
constants though.

* Refactor ASCII color handling into a fixture

* Revert back to using explicit color codes

* In Python 3.5 skip rest of tests that require ordered markup in colored output

Author: nicoddemus

changelog/6658.improvement.rst

@@ -0,0 +1,4 @@
+Code is now highlighted in tracebacks when ``pygments`` is installed.
+
+Users are encouraged to install ``pygments`` into their environment and provide feedback, because
+the plan is to make ``pygments`` a regular dependency in the future.

src/_pytest/_code/code.py

@@ -1039,21 +1039,58 @@ def __init__(
         self.reprfileloc = filelocrepr
         self.style = style
 
+    def _write_entry_lines(self, tw: TerminalWriter) -> None:
+        """Writes the source code portions of a list of traceback entries with syntax highlighting.
+
+        Usually entries are lines like these:
+
+            "     x = 1"
+            ">    assert x == 2"
+            "E    assert 1 == 2"
+
+        This function takes care of rendering the "source" portions of it (the lines without
+        the "E" prefix) using syntax highlighting, taking care to not highlighting the ">"
+        character, as doing so might break line continuations.
+        """
+
+        indent_size = 4
+
+        def is_fail(line):
+            return line.startswith("{}   ".format(FormattedExcinfo.fail_marker))
+
+        if not self.lines:
+            return
+
+        # separate indents and source lines that are not failures: we want to
+        # highlight the code but not the indentation, which may contain markers
+        # such as ">   assert 0"
+        indents = []
+        source_lines = []
+        for line in self.lines:
+            if not is_fail(line):
+                indents.append(line[:indent_size])
+                source_lines.append(line[indent_size:])
+
+        tw._write_source(source_lines, indents)
+
+        # failure lines are always completely red and bold
+        for line in (x for x in self.lines if is_fail(x)):
+            tw.line(line, bold=True, red=True)
+
     def toterminal(self, tw: TerminalWriter) -> None:
         if self.style == "short":
             assert self.reprfileloc is not None
             self.reprfileloc.toterminal(tw)
-            for line in self.lines:
-                red = line.startswith("E   ")
-                tw.line(line, bold=True, red=red)
+            self._write_entry_lines(tw)
             if self.reprlocals:
                 self.reprlocals.toterminal(tw, indent=" " * 8)
             return
+
         if self.reprfuncargs:
             self.reprfuncargs.toterminal(tw)
-        for line in self.lines:
-            red = line.startswith("E   ")
-            tw.line(line, bold=True, red=red)
+
+        self._write_entry_lines(tw)
+
         if self.reprlocals:
             tw.line("")
             self.reprlocals.toterminal(tw)

src/_pytest/_io/__init__.py

@@ -1,3 +1,39 @@
-# Reexport TerminalWriter from here instead of py, to make it easier to
-# extend or swap our own implementation in the future.
-from py.io import TerminalWriter as TerminalWriter  # noqa: F401
+from typing import List
+from typing import Sequence
+
+from py.io import TerminalWriter as BaseTerminalWriter  # noqa: F401
+
+
+class TerminalWriter(BaseTerminalWriter):
+    def _write_source(self, lines: List[str], indents: Sequence[str] = ()) -> None:
+        """Write lines of source code possibly highlighted.
+
+        Keeping this private for now because the API is clunky. We should discuss how
+        to evolve the terminal writer so we can have more precise color support, for example
+        being able to write part of a line in one color and the rest in another, and so on.
+        """
+        if indents and len(indents) != len(lines):
+            raise ValueError(
+                "indents size ({}) should have same size as lines ({})".format(
+                    len(indents), len(lines)
+                )
+            )
+        if not indents:
+            indents = [""] * len(lines)
+        source = "\n".join(lines)
+        new_lines = self._highlight(source).splitlines()
+        for indent, new_line in zip(indents, new_lines):
+            self.line(indent + new_line)
+
+    def _highlight(self, source):
+        """Highlight the given source code according to the "code_highlight" option"""
+        if not self.hasmarkup:
+            return source
+        try:
+            from pygments.formatters.terminal import TerminalFormatter
+            from pygments.lexers.python import PythonLexer
+            from pygments import highlight
+        except ImportError:
+            return source
+        else:
+            return highlight(source, PythonLexer(), TerminalFormatter(bg="dark"))

testing/code/test_terminal_writer.py

@@ -0,0 +1,28 @@
+import re
+from io import StringIO
+
+import pytest
+from _pytest._io import TerminalWriter
+
+
+@pytest.mark.parametrize(
+    "has_markup, expected",
+    [
+        pytest.param(
+            True, "{kw}assert{hl-reset} {number}0{hl-reset}\n", id="with markup"
+        ),
+        pytest.param(False, "assert 0\n", id="no markup"),
+    ],
+)
+def test_code_highlight(has_markup, expected, color_mapping):
+    f = StringIO()
+    tw = TerminalWriter(f)
+    tw.hasmarkup = has_markup
+    tw._write_source(["assert 0"])
+    assert f.getvalue().splitlines(keepends=True) == color_mapping.format([expected])
+
+    with pytest.raises(
+        ValueError,
+        match=re.escape("indents size (2) should have same size as lines (1)"),
+    ):
+        tw._write_source(["assert 0"], [" ", " "])

testing/conftest.py

@@ -1,6 +1,9 @@
+import re
 import sys
+from typing import List
 
 import pytest
+from _pytest.pytester import RunResult
 from _pytest.pytester import Testdir
 
 if sys.gettrace():
@@ -78,6 +81,12 @@ def sep(self, sep, line=None):
         def write(self, msg, **kw):
             self.lines.append((TWMock.WRITE, msg))
 
+        def _write_source(self, lines, indents=()):
+            if not indents:
+                indents = [""] * len(lines)
+            for indent, line in zip(indents, lines):
+                self.line(indent + line)
+
         def line(self, line, **kw):
             self.lines.append(line)
 
@@ -125,3 +134,64 @@ def runtest(self):
 def testdir(testdir: Testdir) -> Testdir:
     testdir.monkeypatch.setenv("PYTEST_DISABLE_PLUGIN_AUTOLOAD", "1")
     return testdir
+
+
+@pytest.fixture(scope="session")
+def color_mapping():
+    """Returns a utility class which can replace keys in strings in the form "{NAME}"
+    by their equivalent ASCII codes in the terminal.
+
+    Used by tests which check the actual colors output by pytest.
+    """
+
+    class ColorMapping:
+        COLORS = {
+            "red": "\x1b[31m",
+            "green": "\x1b[32m",
+            "yellow": "\x1b[33m",
+            "bold": "\x1b[1m",
+            "reset": "\x1b[0m",
+            "kw": "\x1b[94m",
+            "hl-reset": "\x1b[39;49;00m",
+            "function": "\x1b[92m",
+            "number": "\x1b[94m",
+            "str": "\x1b[33m",
+            "print": "\x1b[96m",
+        }
+        RE_COLORS = {k: re.escape(v) for k, v in COLORS.items()}
+
+        @classmethod
+        def format(cls, lines: List[str]) -> List[str]:
+            """Straightforward replacement of color names to their ASCII codes."""
+            return [line.format(**cls.COLORS) for line in lines]
+
+        @classmethod
+        def format_for_fnmatch(cls, lines: List[str]) -> List[str]:
+            """Replace color names for use with LineMatcher.fnmatch_lines"""
+            return [line.format(**cls.COLORS).replace("[", "[[]") for line in lines]
+
+        @classmethod
+        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 requires_ordered_markup(cls, result: RunResult):
+            """Should be called if a test expects markup to appear in the output
+            in the order they were passed, for example:
+
+                tw.write(line, bold=True, red=True)
+
+            In Python 3.5 there's no guarantee that the generated markup will appear
+            in the order called, so we do some limited color testing and skip the rest of
+            the test.
+            """
+            if sys.version_info < (3, 6):
+                # terminal writer.write accepts keyword arguments, so
+                # py36+ is required so the markup appears in the expected order
+                output = result.stdout.str()
+                assert "test session starts" in output
+                assert "\x1b[1m" in output
+                pytest.skip("doing limited testing because lacking ordered markup")
+
+    return ColorMapping