tests

jizongFox · Sep 24, 2021 · e96d889 · e96d889
1 parent 1b49c9f
commit e96d889
Show file tree

Hide file tree

Showing 7 changed files with 174 additions and 42 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -5,6 +5,13 @@ All notable changes to this project will be documented in this file.
 The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
 
+## [10.11.0] - Unreleased
+
+### Added
+
+- Added `suppress` parameter to tracebacks
+- Added `max_frames` parameter to tracebacks
+
 ## [10.10.0] - 2021-09-18
 
 ### Added

diff --git a/docs/source/traceback.rst b/docs/source/traceback.rst
@@ -35,3 +35,42 @@ Rich can be installed as the default traceback handler so that all uncaught exce
     install(show_locals=True)
 
 There are a few options to configure the traceback handler, see :func:`~rich.traceback.install` for details.
+
+
+Suppressing Frames
+------------------
+
+If you are working with a framework (click, django etc), you may only be interested in displaying code in your own application. You can exclude frameworks by setting the `suppress` argument on `Traceback`, `install`, and `print_exception`, which may be a iterable of modules or str paths.
+
+Here's how you would exclude [click](https://click.palletsprojects.com/en/8.0.x/) from Rich exceptions:: 
+
+    import click
+    from rich.traceback import install
+    install(suppress=[click])
+
+Suppressed frames will show the line and file only, without any code.
+
+Max Frames
+----------
+
+A recursion error can generate very large tracebacks that take a while to render and contain a lot of repetitive frames. Rich guards against this with a `max_frames` argument, which defaults to 100. If a traceback contains more than 100 frames then only the first 50, and last 50 will be shown. You can disable this feature by setting `max_frames` to 0.
+
+Here's an example of printing an recursive error::
+
+    from rich.console import Console
+
+
+    def foo(n):
+        return bar(n)
+
+
+    def bar(n):
+        return foo(n)
+
+
+    console = Console()
+
+    try:
+        foo(1)
+    except Exception:
+        console.print_exception(max_frames=20)
diff --git a/examples/recursive_error.py b/examples/recursive_error.py
@@ -0,0 +1,25 @@
+"""
+
+Demonstrates Rich tracebacks for recursion errors.
+
+Rich can exclude frames in the middle to avoid huge tracebacks.
+
+"""
+
+from rich.console import Console
+
+
+def foo(n):
+    return bar(n)
+
+
+def bar(n):
+    return foo(n)
+
+
+console = Console()
+
+try:
+    foo(1)
+except Exception:
+    console.print_exception(max_frames=20)
diff --git a/pyproject.toml b/pyproject.toml
@@ -2,7 +2,7 @@
 name = "rich"
 homepage = "https://github.com/willmcgugan/rich"
 documentation = "https://rich.readthedocs.io/en/latest/"
-version = "10.10.0"
+version = "10.11.0"
 description = "Render rich text, tables, progress bars, syntax highlighting, markdown and more to the terminal"
 authors = ["Will McGugan <willmcgugan@gmail.com>"]
 license = "MIT"

diff --git a/rich/console.py b/rich/console.py
@@ -1705,6 +1705,7 @@ def print_exception(
         word_wrap: bool = False,
         show_locals: bool = False,
         suppress: Iterable[Union[str, ModuleType]] = (),
+        max_frames: int = 100,
     ) -> None:
         """Prints a rich render of the last exception and traceback.
 
@@ -1714,6 +1715,8 @@ def print_exception(
             theme (str, optional): Override pygments theme used in traceback
             word_wrap (bool, optional): Enable word wrapping of long lines. Defaults to False.
             show_locals (bool, optional): Enable display of local variables. Defaults to False.
+            suppress (Iterable[Union[str, ModuleType]]): Optional sequence of modules or paths to exclude from traceback.
+            max_frames (int): Maximum number of frames to show in a traceback, 0 for no maximum. Defaults to 100.
         """
         from .traceback import Traceback
 
@@ -1724,6 +1727,7 @@ def print_exception(
             word_wrap=word_wrap,
             show_locals=show_locals,
             suppress=suppress,
+            max_frames=max_frames,
         )
         self.print(traceback)
 

diff --git a/rich/traceback.py b/rich/traceback.py
@@ -48,6 +48,7 @@ def install(
     show_locals: bool = False,
     indent_guides: bool = True,
     suppress: Iterable[Union[str, ModuleType]] = (),
+    max_frames: int = 100,
 ) -> Callable[[Type[BaseException], BaseException, Optional[TracebackType]], Any]:
     """Install a rich traceback handler.
 
@@ -63,8 +64,7 @@ def install(
         word_wrap (bool, optional): Enable word wrapping of long lines. Defaults to False.
         show_locals (bool, optional): Enable display of local variables. Defaults to False.
         indent_guides (bool, optional): Enable indent guides in code and locals. Defaults to True.
-        suppress (Sequence[Union[str, ModuleType]]): Optional sequence of modules or paths to exclude from traeback.
-
+        suppress (Sequence[Union[str, ModuleType]]): Optional sequence of modules or paths to exclude from traceback.
 
     Returns:
         Callable: The previous exception handler that was replaced.
@@ -89,6 +89,7 @@ def excepthook(
                 show_locals=show_locals,
                 indent_guides=indent_guides,
                 suppress=suppress,
+                max_frames=max_frames,
             )
         )
 
@@ -196,7 +197,8 @@ class Traceback:
         locals_max_length (int, optional): Maximum length of containers before abbreviating, or None for no abbreviation.
             Defaults to 10.
         locals_max_string (int, optional): Maximum length of string before truncating, or None to disable. Defaults to 80.
-        suppress (Sequence[Union[str, ModuleType]]): Optional sequence of modules or paths to exclude from traeback.
+        suppress (Sequence[Union[str, ModuleType]]): Optional sequence of modules or paths to exclude from traceback.
+        max_frames (int): Maximum number of frames to show in a traceback, 0 for no maximum. Defaults to 100.
 
     """
 
@@ -220,6 +222,7 @@ def __init__(
         locals_max_length: int = LOCALS_MAX_LENGTH,
         locals_max_string: int = LOCALS_MAX_STRING,
         suppress: Iterable[Union[str, ModuleType]] = (),
+        max_frames: int = 100,
     ):
         if trace is None:
             exc_type, exc_value, traceback = sys.exc_info()
@@ -248,6 +251,7 @@ def __init__(
                 path = suppress_entity
             path = os.path.normpath(os.path.abspath(path))
             self.suppress.append(path)
+        self.max_frames = max(4, max_frames) if max_frames > 0 else 0
 
     @classmethod
     def from_exception(
@@ -264,6 +268,7 @@ def from_exception(
         locals_max_length: int = LOCALS_MAX_LENGTH,
         locals_max_string: int = LOCALS_MAX_STRING,
         suppress: Iterable[Union[str, ModuleType]] = (),
+        max_frames: int = 100,
     ) -> "Traceback":
         """Create a traceback from exception info
 
@@ -280,7 +285,8 @@ def from_exception(
             locals_max_length (int, optional): Maximum length of containers before abbreviating, or None for no abbreviation.
                 Defaults to 10.
             locals_max_string (int, optional): Maximum length of string before truncating, or None to disable. Defaults to 80.
-            suppress (Iterable[Union[str, ModuleType]]): Optional sequence of modules or paths to exclude from traeback.
+            suppress (Iterable[Union[str, ModuleType]]): Optional sequence of modules or paths to exclude from traceback.
+            max_frames (int): Maximum number of frames to show in a traceback, 0 for no maximum. Defaults to 100.
 
         Returns:
             Traceback: A Traceback instance that may be printed.
@@ -299,6 +305,7 @@ def from_exception(
             locals_max_length=locals_max_length,
             locals_max_string=locals_max_string,
             suppress=suppress,
+            max_frames=max_frames,
         )
 
     @classmethod
@@ -557,7 +564,30 @@ def render_locals(frame: Frame) -> Iterable[ConsoleRenderable]:
                     max_string=self.locals_max_string,
                 )
 
-        for first, frame in loop_first(stack.frames):
+        exclude_frames: Optional[range] = None
+        if self.max_frames != 0:
+            exclude_frames = range(
+                self.max_frames // 2,
+                len(stack.frames) - self.max_frames // 2,
+            )
+
+        excluded = False
+        for frame_index, frame in enumerate(stack.frames):
+
+            if exclude_frames and frame_index in exclude_frames:
+                excluded = True
+                continue
+
+            if excluded:
+                assert exclude_frames is not None
+                yield Text(
+                    f"\n... {len(exclude_frames)} frames hidden ...",
+                    justify="center",
+                    style="traceback.error",
+                )
+                excluded = False
+
+            first = frame_index == 1
             frame_filename = frame.filename
             suppressed = any(frame_filename.startswith(path) for path in self.suppress)
 
@@ -575,43 +605,42 @@ def render_locals(frame: Frame) -> Iterable[ConsoleRenderable]:
             if frame.filename.startswith("<"):
                 yield from render_locals(frame)
                 continue
-            if suppressed:
-                continue
-            try:
-                code = read_code(frame.filename)
-                lexer_name = self._guess_lexer(frame.filename, code)
-                syntax = Syntax(
-                    code,
-                    lexer_name,
-                    theme=theme,
-                    line_numbers=True,
-                    line_range=(
-                        frame.lineno - self.extra_lines,
-                        frame.lineno + self.extra_lines,
-                    ),
-                    highlight_lines={frame.lineno},
-                    word_wrap=self.word_wrap,
-                    code_width=88,
-                    indent_guides=self.indent_guides,
-                    dedent=False,
-                )
-                yield ""
-            except Exception as error:
-                yield Text.assemble(
-                    (f"\n{error}", "traceback.error"),
-                )
-            else:
-                yield (
-                    Columns(
-                        [
-                            syntax,
-                            *render_locals(frame),
-                        ],
-                        padding=1,
+            if not suppressed:
+                try:
+                    code = read_code(frame.filename)
+                    lexer_name = self._guess_lexer(frame.filename, code)
+                    syntax = Syntax(
+                        code,
+                        lexer_name,
+                        theme=theme,
+                        line_numbers=True,
+                        line_range=(
+                            frame.lineno - self.extra_lines,
+                            frame.lineno + self.extra_lines,
+                        ),
+                        highlight_lines={frame.lineno},
+                        word_wrap=self.word_wrap,
+                        code_width=88,
+                        indent_guides=self.indent_guides,
+                        dedent=False,
+                    )
+                    yield ""
+                except Exception as error:
+                    yield Text.assemble(
+                        (f"\n{error}", "traceback.error"),
+                    )
+                else:
+                    yield (
+                        Columns(
+                            [
+                                syntax,
+                                *render_locals(frame),
+                            ],
+                            padding=1,
+                        )
+                        if frame.locals
+                        else syntax
                     )
-                    if frame.locals
-                    else syntax
-                )
 
 
 if __name__ == "__main__":  # pragma: no cover

diff --git a/tests/test_traceback.py b/tests/test_traceback.py
@@ -212,6 +212,34 @@ def test_guess_lexer():
     assert Traceback._guess_lexer("foo", "foo\nbnar") == "text"
 
 
+def test_recursive():
+    def foo(n):
+        return bar(n)
+
+    def bar(n):
+        return foo(n)
+
+    console = Console(width=100, file=io.StringIO())
+    try:
+        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]
+
+
 if __name__ == "__main__":  # pragma: no cover
 
     expected = render(get_exception())