diff --git a/CHANGELOG.md b/CHANGELOG.md index 5001a075c..e4691104b 100644 --- a/CHANGELOG.md +++ b/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 diff --git a/rich/traceback.py b/rich/traceback.py index 5feefb93b..55acaf070 100644 --- a/rich/traceback.py +++ b/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) diff --git a/tests/test_traceback.py b/tests/test_traceback.py index c4994c87b..2830e26aa 100644 --- a/tests/test_traceback.py +++ b/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())