[traceback] Make _rich_traceback_guard more consistent with `_rich_…

…traceback_omit`
Textualize · Apr 26, 2022 · f0ca6ff · f0ca6ff
1 parent 740b89f
commit f0ca6ff
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 3 deletions.
diff --git a/rich/traceback.py b/rich/traceback.py
@@ -367,7 +367,7 @@ 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:
+                if frame_summary.f_locals.get("_rich_traceback_omit", False):
                     continue
                 frame = Frame(
                     filename=filename or "?",
@@ -385,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
@@ -326,7 +326,8 @@ def level1():
         return level2()
 
     def level2():
-        _rich_traceback_omit = rich_traceback_omit_for_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():