Add basic implementation of PEP657 style expression markers in tracebacks (#13102)

* Add very basic implementation of PEP657 style line markers in tracebacks

* Version guard the py3.11 attributes

* Change version guard to make mypy happy. Also, stop using private traceback function

* Version guard unit test

* Fix repr cycle test

* Add changelog entry

* Add pragma no cover for branch where column info is missing

Author: ammaraskar

AUTHORS

@@ -24,6 +24,7 @@ Alice Purcell
 Allan Feldman
 Aly Sivji
 Amir Elkess
+Ammar Askar
 Anatoly Bubenkoff
 Anders Hovmöller
 Andras Mitzki

changelog/10224.improvement.rst

@@ -0,0 +1,18 @@
+pytest's ``short`` and ``long`` traceback styles (:ref:`how-to-modifying-python-tb-printing`)
+now have partial :pep:`657` support and will show specific code segments in the
+traceback.
+
+.. code-block:: pytest
+
+    ================================= FAILURES =================================
+    _______________________ test_gets_correct_tracebacks _______________________
+
+    test_tracebacks.py:12: in test_gets_correct_tracebacks
+        assert manhattan_distance(p1, p2) == 1
+               ^^^^^^^^^^^^^^^^^^^^^^^^^^
+    test_tracebacks.py:6: in manhattan_distance
+        return abs(point_1.x - point_2.x) + abs(point_1.y - point_2.y)
+                               ^^^^^^^^^
+    E   AttributeError: 'NoneType' object has no attribute 'x'
+
+-- by :user:`ammaraskar`

src/_pytest/_code/code.py

@@ -208,6 +208,45 @@ 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]
+
+    # Column and end line numbers introduced in python 3.11
+    if sys.version_info < (3, 11):
+
+        @property
+        def end_lineno_relative(self) -> int | None:
+            return None
+
+        @property
+        def colno(self) -> int | None:
+            return None
+
+        @property
+        def end_colno(self) -> int | None:
+            return None
+    else:
+
+        @property
+        def end_lineno_relative(self) -> int | None:
+            frame_summary = self.get_python_framesummary()
+            if frame_summary.end_lineno is None:  # pragma: no cover
+                return None
+            return frame_summary.end_lineno - 1 - self.frame.code.firstlineno
+
+        @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 +895,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 +911,74 @@ 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 = _byte_offset_to_character_offset(raw_line, colno)
+        end_char_offset = _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 +1038,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()
+                line_index = entry.relline
+                end_line_index = entry.end_lineno_relative
+                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}"
@@ -1374,6 +1485,12 @@ def getfslineno(obj: object) -> tuple[str | Path, int]:
     return code.path, code.firstlineno
 
 
+def _byte_offset_to_character_offset(str, offset):
+    """Converts a byte based offset in a string to a code-point."""
+    as_utf8 = str.encode("utf-8")
+    return len(as_utf8[:offset].decode("utf-8", errors="replace"))
+
+
 # Relative paths that we use to filter traceback entries from appearing to the user;
 # see filter_traceback.
 # note: if we need to add more paths than what we have now we should probably use a list

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

@@ -883,6 +883,37 @@ def entry():
         assert basename in str(reprtb.reprfileloc.path)
         assert reprtb.reprfileloc.lineno == 3
 
+    @pytest.mark.skipif(
+        "sys.version_info < (3,11)",
+        reason="Column level traceback info added in python 3.11",
+    )
+    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(
             """
@@ -1343,7 +1374,7 @@ def g():
                 raise ValueError()
 
             def h():
-                raise AttributeError()
+                if True: raise AttributeError()
         """
         )
         excinfo = pytest.raises(AttributeError, mod.f)
@@ -1404,12 +1435,22 @@ 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 line.endswith("mod.py")
-        assert tw_mock.lines[47] == ":15: AttributeError"
+        # On python 3.11 and greater, check for carets in the traceback.
+        if sys.version_info >= (3, 11):
+            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[48] == ":15: AttributeError"
+        else:
+            assert tw_mock.lines[43] == ">       if True: raise AttributeError()"
+            assert tw_mock.lines[44] == "E       AttributeError"
+            assert tw_mock.lines[45] == ""
+            line = tw_mock.get_write_msg(46)
+            assert line.endswith("mod.py")
+            assert tw_mock.lines[47] == ":15: AttributeError"
 
     @pytest.mark.parametrize("mode", ["from_none", "explicit_suppress"])
     def test_exc_repr_chain_suppression(self, importasmod, mode, tw_mock):
@@ -1528,23 +1569,44 @@ def unreraise():
         r = excinfo.getrepr(style="short")
         r.toterminal(tw_mock)
         out = "\n".join(line for line in tw_mock.lines if isinstance(line, str))
-        expected_out = textwrap.dedent(
-            """\
-            :13: in unreraise
-                reraise()
-            :10: in reraise
-                raise Err() from e
-            E   test_exc_chain_repr_cycle0.mod.Err
-
-            During handling of the above exception, another exception occurred:
-            :15: in unreraise
-                raise e.__cause__
-            :8: in reraise
-                fail()
-            :5: in fail
-                return 0 / 0
-            E   ZeroDivisionError: division by zero"""
-        )
+        # Assert highlighting carets in python3.11+
+        if sys.version_info >= (3, 11):
+            expected_out = textwrap.dedent(
+                """\
+                :13: in unreraise
+                    reraise()
+                :10: in reraise
+                    raise Err() from e
+                E   test_exc_chain_repr_cycle0.mod.Err
+
+                During handling of the above exception, another exception occurred:
+                :15: in unreraise
+                    raise e.__cause__
+                :8: in reraise
+                    fail()
+                :5: in fail
+                    return 0 / 0
+                           ^^^^^
+                E   ZeroDivisionError: division by zero"""
+            )
+        else:
+            expected_out = textwrap.dedent(
+                """\
+                :13: in unreraise
+                    reraise()
+                :10: in reraise
+                    raise Err() from e
+                E   test_exc_chain_repr_cycle0.mod.Err
+
+                During handling of the above exception, another exception occurred:
+                :15: in unreraise
+                    raise e.__cause__
+                :8: in reraise
+                    fail()
+                :5: in fail
+                    return 0 / 0
+                E   ZeroDivisionError: division by zero"""
+            )
         assert out == expected_out
 
     def test_exec_type_error_filter(self, importasmod):