Merge pull request #2226 from Textualize/feat-add-traceback-frame-opt-out-machinery

[traceback] Add traceback frame opt-out via a `_rich_traceback_omit = True` machinery

Author: willmcgugan

CHANGELOG.md

@@ -11,6 +11,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 - Ability to change terminal window title https://github.com/Textualize/rich/pull/2200
 - Added show_speed parameter to progress.track which will show the speed when the total is not known
+- Python blocks can now opt out from being rendered in tracebacks's frames, by setting a `_rich_traceback_omit = True` in their local scope https://github.com/Textualize/rich/issues/2207
 
 ### Fixed
 

rich/traceback.py

@@ -367,6 +367,8 @@ def safe_str(_object: Any) -> str:
                 if filename and not filename.startswith("<"):
                     if not os.path.isabs(filename):
                         filename = os.path.join(_IMPORT_CWD, filename)
+                if frame_summary.f_locals.get("_rich_traceback_omit", False):
+                    continue
                 frame = Frame(
                     filename=filename or "?",
                     lineno=line_no,
@@ -383,7 +385,7 @@ def safe_str(_object: Any) -> str:
                     else None,
                 )
                 append(frame)
-                if "_rich_traceback_guard" in frame_summary.f_locals:
+                if frame_summary.f_locals.get("_rich_traceback_guard", False):
                     del stack.frames[:]
 
             cause = getattr(exc_value, "__cause__", None)

tests/test_traceback.py

@@ -1,6 +1,7 @@
 import io
 import re
 import sys
+from typing import List
 
 import pytest
 
@@ -307,6 +308,42 @@ def test_suppress():
         assert "foo" in traceback.suppress[1]
 
 
+@pytest.mark.parametrize(
+    "rich_traceback_omit_for_level2,expected_frames_length,expected_frame_names",
+    (
+        # fmt: off
+        [True, 3, ["test_rich_traceback_omit_optional_local_flag", "level1", "level3"]],
+        [False, 4, ["test_rich_traceback_omit_optional_local_flag", "level1", "level2", "level3"]],
+        # fmt: on
+    ),
+)
+def test_rich_traceback_omit_optional_local_flag(
+    rich_traceback_omit_for_level2: bool,
+    expected_frames_length: int,
+    expected_frame_names: List[str],
+):
+    def level1():
+        return level2()
+
+    def level2():
+        # true-ish values are enough to trigger the opt-out:
+        _rich_traceback_omit = 1 if rich_traceback_omit_for_level2 else 0
+        return level3()
+
+    def level3():
+        return 1 / 0
+
+    try:
+        level1()
+    except Exception:
+        exc_type, exc_value, traceback = sys.exc_info()
+        trace = Traceback.from_exception(exc_type, exc_value, traceback).trace
+        frames = trace.stacks[0].frames
+        assert len(frames) == expected_frames_length
+        frame_names = [f.name for f in frames]
+        assert frame_names == expected_frame_names
+
+
 if __name__ == "__main__":  # pragma: no cover
 
     expected = render(get_exception())