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

This commit is contained in:
Olivier Philippon 2022-04-26 12:40:51 +01:00
parent 49840e3b82
commit 740b89f49a
3 changed files with 39 additions and 0 deletions

View file

@ -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

View file

@ -367,6 +367,8 @@ class Traceback:
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) is True:
continue
frame = Frame(
filename=filename or "?",
lineno=line_no,

View file

@ -1,6 +1,7 @@
import io
import re
import sys
from typing import List
import pytest
@ -307,6 +308,41 @@ 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():
_rich_traceback_omit = rich_traceback_omit_for_level2
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())