Repository: Textualize/rich
File: tests/test_traceback.py (Python)

Code before:
        foo(1)
    except Exception:
        console.print_exception(max_frames=6)
    result = console.file.getvalue()
    print(result)
    assert "frames hidden" in result
    assert result.count("in foo") < 4


def test_suppress():
    try:
        1 / 0
    except Exception:
        traceback = Traceback(suppress=[pytest, "foo"])
        assert len(traceback.suppress) == 2
        assert "pytest" in traceback.suppress[0]
        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]

Review comment (suggestion):
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 :slightly_smiling_face: 

Code after:
        foo(1)
    except Exception:
        console.print_exception(max_frames=6)
    result = console.file.getvalue()
    print(result)
    assert "frames hidden" in result
    assert result.count("in foo") < 4


def test_suppress():
    try:
        1 / 0
    except Exception:
        traceback = Traceback(suppress=[pytest, "foo"])
        assert len(traceback.suppress) == 2
        assert "pytest" in traceback.suppress[0]
        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