Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

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

Merged
merged 2 commits into from Apr 26, 2022
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Jump to
Jump to file
Failed to load files.
Diff view
Diff view
1 change: 1 addition & 0 deletions CHANGELOG.md
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
4 changes: 3 additions & 1 deletion rich/traceback.py
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):
continue
frame = Frame(
filename=filename or "?",
lineno=line_no,
Expand All @@ -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)
Expand Down
37 changes: 37 additions & 0 deletions tests/test_traceback.py
@@ -1,6 +1,7 @@
import io
import re
import sys
from typing import List

import pytest

Expand Down Expand Up @@ -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
Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I added these # fmt: annotations just to avoid Black changing the second test case so that it spawns across multiple lines when the 1st one doesn't, as they are functionally the same thing
But happy to let Black format everything and remove those if you prefer, of course 🙂

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I like this. Black completely ruins the readability of some parametrised tests

Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Fine by me!

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