Skip to content

Commit

Permalink
[traceback] Add traceback frame opt-out via a `_rich_traceback_omit =…
Browse files Browse the repository at this point in the history
… True` machinery
  • Loading branch information
Olivier Philippon committed Apr 26, 2022
1 parent 49840e3 commit 740b89f
Show file tree
Hide file tree
Showing 3 changed files with 39 additions and 0 deletions.
1 change: 1 addition & 0 deletions CHANGELOG.md
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down
2 changes: 2 additions & 0 deletions rich/traceback.py
Original file line number Diff line number Diff line change
Expand Up @@ -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) is True:
continue
frame = Frame(
filename=filename or "?",
lineno=line_no,
Expand Down
36 changes: 36 additions & 0 deletions tests/test_traceback.py
Original file line number Diff line number Diff line change
@@ -1,6 +1,7 @@
import io
import re
import sys
from typing import List

import pytest

Expand Down Expand Up @@ -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())
Expand Down

0 comments on commit 740b89f

Please sign in to comment.