Add very basic implementation of PEP657 style line markers in tracebacks

Author: ammaraskar

src/_pytest/_code/code.py

@@ -208,6 +208,26 @@ def with_repr_style(
     def lineno(self) -> int:
         return self._rawentry.tb_lineno - 1
 
+    def get_python_framesummary(self) -> traceback.FrameSummary:
+        # Python's built-in traceback module implements all the nitty gritty
+        # details to get column numbers of out frames.
+        stack_summary = traceback.extract_tb(self._rawentry, limit=1)
+        return stack_summary[0]
+
+    @property
+    def end_lineno(self) -> int:
+        return self.get_python_framesummary().end_lineno - 1
+
+    @property
+    def colno(self) -> int | None:
+        """Starting byte offset of the expression in the traceback entry."""
+        return self.get_python_framesummary().colno
+
+    @property
+    def end_colno(self) -> int | None:
+        """Ending byte offset of the expression in the traceback entry."""
+        return self.get_python_framesummary().end_colno
+
     @property
     def frame(self) -> Frame:
         return Frame(self._rawentry.tb_frame)
@@ -856,6 +876,9 @@ def get_source(
         line_index: int = -1,
         excinfo: ExceptionInfo[BaseException] | None = None,
         short: bool = False,
+        end_line_index: int | None = None,
+        colno: int | None = None,
+        end_colno: int | None = None,
     ) -> list[str]:
         """Return formatted and marked up source lines."""
         lines = []
@@ -869,17 +892,76 @@ def get_source(
         space_prefix = "    "
         if short:
             lines.append(space_prefix + source.lines[line_index].strip())
+            lines.extend(
+                self.get_highlight_arrows_for_line(
+                    raw_line=source.raw_lines[line_index],
+                    line=source.lines[line_index].strip(),
+                    lineno=line_index,
+                    end_lineno=end_line_index,
+                    colno=colno,
+                    end_colno=end_colno,
+                )
+            )
         else:
             for line in source.lines[:line_index]:
                 lines.append(space_prefix + line)
             lines.append(self.flow_marker + "   " + source.lines[line_index])
+            lines.extend(
+                self.get_highlight_arrows_for_line(
+                    raw_line=source.raw_lines[line_index],
+                    line=source.lines[line_index],
+                    lineno=line_index,
+                    end_lineno=end_line_index,
+                    colno=colno,
+                    end_colno=end_colno,
+                )
+            )
             for line in source.lines[line_index + 1 :]:
                 lines.append(space_prefix + line)
         if excinfo is not None:
             indent = 4 if short else self._getindent(source)
             lines.extend(self.get_exconly(excinfo, indent=indent, markall=True))
         return lines
 
+    def get_highlight_arrows_for_line(
+        self,
+        line: str,
+        raw_line: str,
+        lineno: int | None,
+        end_lineno: int | None,
+        colno: int | None,
+        end_colno: int | None,
+    ) -> list[str]:
+        """Return characters highlighting a source line.
+
+        Example with colno and end_colno pointing to the bar expression:
+                   "foo() + bar()"
+        returns    "        ^^^^^"
+        """
+        if lineno != end_lineno:
+            # Don't handle expressions that span multiple lines.
+            return []
+        if colno is None or end_colno is None:
+            # Can't do anything without column information.
+            return []
+
+        num_stripped_chars = len(raw_line) - len(line)
+
+        start_char_offset = traceback._byte_offset_to_character_offset(raw_line, colno)
+        end_char_offset = traceback._byte_offset_to_character_offset(
+            raw_line, end_colno
+        )
+        num_carets = end_char_offset - start_char_offset
+        # If the highlight would span the whole line, it is redundant, don't
+        # show it.
+        if num_carets >= len(line.strip()):
+            return []
+
+        highlights = "    "
+        highlights += " " * (start_char_offset - num_stripped_chars + 1)
+        highlights += "^" * num_carets
+        return [highlights]
+
     def get_exconly(
         self,
         excinfo: ExceptionInfo[BaseException],
@@ -939,11 +1021,23 @@ def repr_traceback_entry(
             if source is None:
                 source = Source("???")
                 line_index = 0
+                end_line_index, colno, end_colno = None, None, None
             else:
                 line_index = entry.lineno - entry.getfirstlinesource()
+                end_line_index = entry.end_lineno - entry.getfirstlinesource()
+                colno = entry.colno
+                end_colno = entry.end_colno
             short = style == "short"
             reprargs = self.repr_args(entry) if not short else None
-            s = self.get_source(source, line_index, excinfo, short=short)
+            s = self.get_source(
+                source=source,
+                line_index=line_index,
+                excinfo=excinfo,
+                short=short,
+                end_line_index=end_line_index,
+                colno=colno,
+                end_colno=end_colno,
+            )
             lines.extend(s)
             if short:
                 message = f"in {entry.name}"

src/_pytest/_code/source.py

@@ -22,19 +22,24 @@ class Source:
     def __init__(self, obj: object = None) -> None:
         if not obj:
             self.lines: list[str] = []
+            self.raw_lines: list[str] = []
         elif isinstance(obj, Source):
             self.lines = obj.lines
+            self.raw_lines = obj.raw_lines
         elif isinstance(obj, (tuple, list)):
             self.lines = deindent(x.rstrip("\n") for x in obj)
+            self.raw_lines = list(x.rstrip("\n") for x in obj)
         elif isinstance(obj, str):
             self.lines = deindent(obj.split("\n"))
+            self.raw_lines = obj.split("\n")
         else:
             try:
                 rawcode = getrawcode(obj)
                 src = inspect.getsource(rawcode)
             except TypeError:
                 src = inspect.getsource(obj)  # type: ignore[arg-type]
             self.lines = deindent(src.split("\n"))
+            self.raw_lines = src.split("\n")
 
     def __eq__(self, other: object) -> bool:
         if not isinstance(other, Source):
@@ -58,6 +63,7 @@ def __getitem__(self, key: int | slice) -> str | Source:
                 raise IndexError("cannot slice a Source with a step")
             newsource = Source()
             newsource.lines = self.lines[key.start : key.stop]
+            newsource.raw_lines = self.raw_lines[key.start : key.stop]
             return newsource
 
     def __iter__(self) -> Iterator[str]:
@@ -74,13 +80,15 @@ def strip(self) -> Source:
         while end > start and not self.lines[end - 1].strip():
             end -= 1
         source = Source()
+        source.raw_lines = self.raw_lines
         source.lines[:] = self.lines[start:end]
         return source
 
     def indent(self, indent: str = " " * 4) -> Source:
         """Return a copy of the source object with all lines indented by the
         given indent-string."""
         newsource = Source()
+        newsource.raw_lines = self.raw_lines
         newsource.lines = [(indent + line) for line in self.lines]
         return newsource
 
@@ -102,6 +110,7 @@ def deindent(self) -> Source:
         """Return a new Source object deindented."""
         newsource = Source()
         newsource.lines[:] = deindent(self.lines)
+        newsource.raw_lines = self.raw_lines
         return newsource
 
     def __str__(self) -> str:
@@ -120,6 +129,7 @@ def findsource(obj) -> tuple[Source | None, int]:
         return None, -1
     source = Source()
     source.lines = [line.rstrip() for line in sourcelines]
+    source.raw_lines = sourcelines
     return source, lineno
 
 

testing/code/test_excinfo.py

@@ -849,6 +849,33 @@ def entry():
         assert basename in str(reprtb.reprfileloc.path)
         assert reprtb.reprfileloc.lineno == 3
 
+    def test_repr_traceback_entry_short_carets(self, importasmod) -> None:
+        mod = importasmod(
+            """
+            def div_by_zero():
+                return 1 / 0
+            def func1():
+                return 42 + div_by_zero()
+            def entry():
+                func1()
+        """
+        )
+        excinfo = pytest.raises(ZeroDivisionError, mod.entry)
+        p = FormattedExcinfo(style="short")
+        reprtb = p.repr_traceback_entry(excinfo.traceback[-3])
+        assert len(reprtb.lines) == 1
+        assert reprtb.lines[0] == "    func1()"
+
+        reprtb = p.repr_traceback_entry(excinfo.traceback[-2])
+        assert len(reprtb.lines) == 2
+        assert reprtb.lines[0] == "    return 42 + div_by_zero()"
+        assert reprtb.lines[1] == "                ^^^^^^^^^^^^^"
+
+        reprtb = p.repr_traceback_entry(excinfo.traceback[-1])
+        assert len(reprtb.lines) == 2
+        assert reprtb.lines[0] == "    return 1 / 0"
+        assert reprtb.lines[1] == "           ^^^^^"
+
     def test_repr_tracebackentry_no(self, importasmod):
         mod = importasmod(
             """
@@ -1309,7 +1336,7 @@ def g():
                 raise ValueError()
 
             def h():
-                raise AttributeError()
+                if True: raise AttributeError()
         """
         )
         excinfo = pytest.raises(AttributeError, mod.f)
@@ -1370,12 +1397,13 @@ def h():
         assert tw_mock.lines[40] == ("_ ", None)
         assert tw_mock.lines[41] == ""
         assert tw_mock.lines[42] == "    def h():"
-        assert tw_mock.lines[43] == ">       raise AttributeError()"
-        assert tw_mock.lines[44] == "E       AttributeError"
-        assert tw_mock.lines[45] == ""
-        line = tw_mock.get_write_msg(46)
+        assert tw_mock.lines[43] == ">       if True: raise AttributeError()"
+        assert tw_mock.lines[44] == "                 ^^^^^^^^^^^^^^^^^^^^^^"
+        assert tw_mock.lines[45] == "E       AttributeError"
+        assert tw_mock.lines[46] == ""
+        line = tw_mock.get_write_msg(47)
         assert line.endswith("mod.py")
-        assert tw_mock.lines[47] == ":15: AttributeError"
+        assert tw_mock.lines[48] == ":15: AttributeError"
 
     @pytest.mark.parametrize("mode", ["from_none", "explicit_suppress"])
     def test_exc_repr_chain_suppression(self, importasmod, mode, tw_mock):
@@ -1509,6 +1537,7 @@ def unreraise():
                 fail()
             :5: in fail
                 return 0 / 0
+                       ^^^^^
             E   ZeroDivisionError: division by zero"""
         )
         assert out == expected_out