Parse on ReprEntry

Author: nicoddemus

src/_pytest/_code/code.py

@@ -29,7 +29,6 @@
 import py
 
 import _pytest
-from _pytest._io import ParsedTraceback
 from _pytest._io import TerminalWriter
 from _pytest._io.saferepr import safeformat
 from _pytest._io.saferepr import saferepr
@@ -1039,26 +1038,31 @@ def __init__(
         self.style = style
 
     def toterminal(self, tw: TerminalWriter) -> None:
+        def is_fail(line):
+            return line.startswith("{}   ".format(FormattedExcinfo.fail_marker))
 
-        traceback_lines = ParsedTraceback.from_lines(self.lines)
-        source = "\n".join(traceback_lines.sources + traceback_lines.errors) + "\n"
+        indents_and_source_lines = [
+            (x[:4], x[4:]) for x in self.lines if not is_fail(x)
+        ]
+        indents, source_lines = list(zip(*indents_and_source_lines))
+
+        def write_tb_lines():
+            tw._write_source(source_lines, indents)
+            for line in (x for x in self.lines if is_fail(x)):
+                tw.line(line, bold=True, red=True)
 
         if self.style == "short":
             assert self.reprfileloc is not None
             self.reprfileloc.toterminal(tw)
-            tw.write_source(source)
-            for line in traceback_lines.explanations:
-                tw.line(line, bold=True, red=True)
+            write_tb_lines()
             if self.reprlocals:
                 self.reprlocals.toterminal(tw, indent=" " * 8)
             return
 
         if self.reprfuncargs:
             self.reprfuncargs.toterminal(tw)
 
-        tw.write_source(source)
-        for line in traceback_lines.explanations:
-            tw.line(line, bold=True, red=True)
+        write_tb_lines()
 
         if self.reprlocals:
             tw.line("")

src/_pytest/_io/__init__.py

@@ -1,16 +1,24 @@
-from enum import auto
-from enum import Enum
+from typing import List
 from typing import Sequence
-from typing import Tuple
 
-import attr
 from py.io import TerminalWriter as BaseTerminalWriter  # noqa: F401
 
 
 class TerminalWriter(BaseTerminalWriter):
-    def write_source(self, source: str) -> None:
+    def _write_source(self, lines: List[str], indents: Sequence[str] = ()) -> None:
         """Write lines of source code possibly highlighted."""
-        self.write(self._highlight(source))
+        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"""
@@ -24,44 +32,3 @@ def _highlight(self, source):
             return source
         else:
             return highlight(source, PythonLexer(), TerminalFormatter(bg="dark"))
-
-
-@attr.s(frozen=True)
-class ParsedTraceback:
-    sources = attr.ib()  # type: Tuple[str, ...]
-    errors = attr.ib()  # type: Tuple[str, ...]
-    explanations = attr.ib()  # type: Tuple[str, ...]
-
-    ERROR_PREFIX = ">   "
-    EXPLANATION_PREFIX = "E   "
-
-    class State(Enum):
-        source = auto()
-        error = auto()
-        explanation = auto()
-
-    @classmethod
-    def from_lines(cls, lines: Sequence[str]) -> "ParsedTraceback":
-        sources = []
-        errors = []
-        explanations = []
-
-        state = cls.State.source
-
-        for line in lines:
-            if state == cls.State.source and line.startswith(cls.ERROR_PREFIX):
-                state = cls.State.error
-            if line.startswith(cls.EXPLANATION_PREFIX):
-                state = cls.State.explanation
-
-            if state == cls.State.source:
-                sources.append(line)
-            elif state == cls.State.error:
-                errors.append(line)
-            else:
-                assert state == cls.State.explanation, "unknown state {!r}".format(
-                    state
-                )
-                explanations.append(line)
-
-        return ParsedTraceback(tuple(sources), tuple(errors), tuple(explanations))

testing/code/test_code.py

@@ -7,7 +7,6 @@
 from _pytest._code import ExceptionInfo
 from _pytest._code import Frame
 from _pytest._code.code import ReprFuncArgs
-from _pytest._io import ParsedTraceback
 
 
 def test_ne() -> None:
@@ -181,59 +180,3 @@ def test_not_raise_exception_with_mixed_encoding(self, tw_mock) -> None:
             tw_mock.lines[0]
             == r"unicode_string = São Paulo, utf8_string = b'S\xc3\xa3o Paulo'"
         )
-
-
-class TestParsedTraceback:
-    def test_simple(self):
-        lines = [
-            "    def test():\n",
-            "        print('hello')\n",
-            "        foo = y()\n",
-            ">       assert foo == z()\n",
-            "E       assert 1 == 2\n",
-            "E        +  where 2 = z()\n",
-        ]
-        traceback_lines = ParsedTraceback.from_lines(lines)
-        assert traceback_lines.sources == (
-            "    def test():\n",
-            "        print('hello')\n",
-            "        foo = y()\n",
-        )
-        assert traceback_lines.errors == (">       assert foo == z()\n",)
-        assert traceback_lines.explanations == (
-            "E       assert 1 == 2\n",
-            "E        +  where 2 = z()\n",
-        )
-
-    def test_with_continuations(self):
-        lines = [
-            "    def test():\n",
-            "        foo = \\\n",
-            "            y()\n",
-            ">       assert foo == \\\n",
-            "               z()\n",
-            "E       assert 1 == 2\n",
-            "E        +  where 2 = z()\n",
-        ]
-        result = ParsedTraceback.from_lines(lines)
-        assert result.sources == (
-            "    def test():\n",
-            "        foo = \\\n",
-            "            y()\n",
-        )
-        assert result.errors == (">       assert foo == \\\n", "               z()\n",)
-        assert result.explanations == (
-            "E       assert 1 == 2\n",
-            "E        +  where 2 = z()\n",
-        )
-
-    def test_no_flow_lines(self):
-        lines = [
-            "    assert 0\n",
-            "E   assert 0\n",
-        ]
-
-        result = ParsedTraceback.from_lines(lines)
-        assert result.sources == ("    assert 0\n",)
-        assert result.errors == ()
-        assert result.explanations == ("E   assert 0\n",)

testing/code/test_terminal_writer.py

@@ -0,0 +1,45 @@
+import re
+from io import StringIO
+
+import pytest
+from _pytest._io import TerminalWriter
+
+
+# TODO: move this and the other two related attributes from test_terminal.py into conftest as a
+# fixture
+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",
+}
+
+
+@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):
+    f = StringIO()
+    tw = TerminalWriter(f)
+    tw.hasmarkup = has_markup
+    tw._write_source(["assert 0"])
+    assert f.getvalue() == expected.format(**COLORS)
+
+    with pytest.raises(
+        ValueError,
+        match=re.escape("indents size (2) should have same size as lines (1)"),
+    ):
+        tw._write_source(["assert 0"], [" ", " "])

testing/test_terminal.py

@@ -34,6 +34,8 @@
     "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()}
 
@@ -2025,21 +2027,43 @@ def test_via_exec(testdir: Testdir) -> None:
     )
 
 
-def test_code_highlight(testdir: Testdir) -> None:
-    testdir.makepyfile(
+class TestCodeHighlight:
+    def test_code_highlight_simple(self, testdir: Testdir) -> None:
+        testdir.makepyfile(
+            """
+            def test_foo():
+                assert 1 == 10
         """
-        def test_foo():
-            assert 1 == 10
-    """
-    )
-    result = testdir.runpytest("--color=yes")
-    result.stdout.fnmatch_lines(
-        [
-            line.format(**COLORS).replace("[", "[[]")
-            for line in [
-                "    {kw}def{hl-reset} {function}test_foo{hl-reset}():",
-                ">       {kw}assert{hl-reset} {number}1{hl-reset} == {number}10{hl-reset}",
-                "{bold}{red}E       assert 1 == 10{reset}",
+        )
+        result = testdir.runpytest("--color=yes")
+        result.stdout.fnmatch_lines(
+            [
+                line.format(**COLORS).replace("[", "[[]")
+                for line in [
+                    "    {kw}def{hl-reset} {function}test_foo{hl-reset}():",
+                    ">       {kw}assert{hl-reset} {number}1{hl-reset} == {number}10{hl-reset}",
+                    "{bold}{red}E       assert 1 == 10{reset}",
+                ]
             ]
-        ]
-    )
+        )
+
+    def test_code_highlight_continuation(self, testdir: Testdir) -> None:
+        testdir.makepyfile(
+            """
+            def test_foo():
+                print('''
+                '''); assert 0
+        """
+        )
+        result = testdir.runpytest("--color=yes")
+        result.stdout.fnmatch_lines(
+            [
+                line.format(**COLORS).replace("[", "[[]")
+                for line in [
+                    "    {kw}def{hl-reset} {function}test_foo{hl-reset}():",
+                    "        {print}print{hl-reset}({str}'''{hl-reset}{str}{hl-reset}",
+                    ">   {str}    {hl-reset}{str}'''{hl-reset}); {kw}assert{hl-reset} {number}0{hl-reset}",
+                    "{bold}{red}E       assert 0{reset}",
+                ]
+            ]
+        )