Use code highlighting if pygments is installed

Author: nicoddemus

src/_pytest/_code/code.py

@@ -1038,20 +1038,27 @@ def __init__(
         self.style = style
 
     def toterminal(self, tw: TerminalWriter) -> None:
+        error_prefix = "E   "
+        src = "\n".join(x for x in self.lines if not x.startswith(error_prefix))
+        error_lines = (x for x in self.lines if x.startswith(error_prefix))
+
         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)
+            tw.write_source(src + "\n")
+            for line in error_lines:
+                tw.line(line, red=True, bold=True)
             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)
+
+        tw.write_source(src + "\n")
+        for line in error_lines:
+            tw.line(line, red=True, bold=True)
+
         if self.reprlocals:
             tw.line("")
             self.reprlocals.toterminal(tw)

src/_pytest/_io/__init__.py

@@ -1,3 +1,22 @@
 # 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 py.io import TerminalWriter as BaseTerminalWriter  # noqa: F401
+
+
+class TerminalWriter(BaseTerminalWriter):
+    def write_source(self, source: str) -> None:
+        """Write lines of source code possibly highlighted."""
+        self.write(self._highlight(source))
+
+    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())

testing/test_terminal.py

@@ -916,16 +916,9 @@ def test_this():
                 "=*= FAILURES =*=",
                 "{red}{bold}_*_ test_this _*_{reset}",
                 "",
-                "{bold}    def test_this():{reset}",
-                "{bold}>       fail(){reset}",
-                "",
                 "{bold}{red}test_color_yes.py{reset}:5: ",
                 "_ _ * _ _*",
                 "",
-                "{bold}    def fail():{reset}",
-                "{bold}>       assert 0{reset}",
-                "{bold}{red}E       assert 0{reset}",
-                "",
                 "{bold}{red}test_color_yes.py{reset}:2: AssertionError",
                 "{red}=*= {red}{bold}1 failed{reset}{red} in *s{reset}{red} =*={reset}",
             ]
@@ -944,10 +937,7 @@ def test_this():
                 "=*= FAILURES =*=",
                 "{red}{bold}_*_ test_this _*_{reset}",
                 "{bold}{red}test_color_yes.py{reset}:5: in test_this",
-                "{bold}    fail(){reset}",
                 "{bold}{red}test_color_yes.py{reset}:2: in fail",
-                "{bold}    assert 0{reset}",
-                "{bold}{red}E   assert 0{reset}",
                 "{red}=*= {red}{bold}1 failed{reset}{red} in *s{reset}{red} =*={reset}",
             ]
         ]
@@ -2019,3 +2009,19 @@ def test_via_exec(testdir: Testdir) -> None:
     result.stdout.fnmatch_lines(
         ["test_via_exec.py::test_via_exec <- <string> PASSED*", "*= 1 passed in *"]
     )
+
+
+def test_code_highlight(testdir: Testdir) -> None:
+    testdir.makepyfile(
+        """
+        def test_foo():
+            assert 1 == 10
+    """
+    )
+    result = testdir.runpytest("--color=yes")
+    result.stdout.fnmatch_lines(
+        [
+            "*\x1b[34mdef\x1b[39;49;00m \x1b[32mtest_foo\x1b[39;49;00m():",
+            "*E       assert 1 == 10*",
+        ]
+    )

tox.ini

@@ -49,6 +49,7 @@ deps =
     numpy: numpy
     pexpect: pexpect
     pluggymaster: git+https://github.com/pytest-dev/pluggy.git@master
+    pygments
     twisted: twisted
     xdist: pytest-xdist>=1.13
     {env:_PYTEST_TOX_EXTRA_DEP:}