Skip to content

Commit

Permalink
add suppress option to tracebacks
Browse files Browse the repository at this point in the history
  • Loading branch information
willmcgugan committed Sep 23, 2021
1 parent 522f841 commit 1b49c9f
Show file tree
Hide file tree
Showing 3 changed files with 54 additions and 4 deletions.
23 changes: 23 additions & 0 deletions examples/suppress.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,23 @@
try:
import click
except ImportError:
print("Please install click for this example")
print(" pip install click")
exit()

from rich.traceback import install

install(suppress=[click])


@click.command()
@click.option("--count", default=1, help="Number of greetings.")
def hello(count):
"""Simple program that greets NAME for a total of COUNT times."""
1 / 0
for x in range(count):
click.echo(f"Hello {name}!")


if __name__ == "__main__":
hello()
4 changes: 3 additions & 1 deletion rich/console.py
Original file line number Diff line number Diff line change
Expand Up @@ -13,7 +13,7 @@
from inspect import isclass
from itertools import islice
from time import monotonic
from types import FrameType, TracebackType
from types import FrameType, TracebackType, ModuleType
from typing import (
IO,
TYPE_CHECKING,
Expand Down Expand Up @@ -1704,6 +1704,7 @@ def print_exception(
theme: Optional[str] = None,
word_wrap: bool = False,
show_locals: bool = False,
suppress: Iterable[Union[str, ModuleType]] = (),
) -> None:
"""Prints a rich render of the last exception and traceback.
Expand All @@ -1722,6 +1723,7 @@ def print_exception(
theme=theme,
word_wrap=word_wrap,
show_locals=show_locals,
suppress=suppress,
)
self.print(traceback)

Expand Down
31 changes: 28 additions & 3 deletions rich/traceback.py
Original file line number Diff line number Diff line change
Expand Up @@ -5,8 +5,8 @@
import sys
from dataclasses import dataclass, field
from traceback import walk_tb
from types import TracebackType
from typing import Any, Callable, Dict, Iterable, List, Optional, Type
from types import ModuleType, TracebackType
from typing import Any, Callable, Dict, Iterable, List, Optional, Sequence, Type, Union

from pygments.lexers import guess_lexer_for_filename
from pygments.token import Comment, Keyword, Name, Number, Operator, String
Expand Down Expand Up @@ -47,6 +47,7 @@ def install(
word_wrap: bool = False,
show_locals: bool = False,
indent_guides: bool = True,
suppress: Iterable[Union[str, ModuleType]] = (),
) -> Callable[[Type[BaseException], BaseException, Optional[TracebackType]], Any]:
"""Install a rich traceback handler.
Expand All @@ -62,6 +63,8 @@ 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.
Returns:
Callable: The previous exception handler that was replaced.
Expand All @@ -85,6 +88,7 @@ def excepthook(
word_wrap=word_wrap,
show_locals=show_locals,
indent_guides=indent_guides,
suppress=suppress,
)
)

Expand Down Expand Up @@ -192,6 +196,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.
"""

LEXERS = {
Expand All @@ -213,6 +219,7 @@ def __init__(
indent_guides: bool = True,
locals_max_length: int = LOCALS_MAX_LENGTH,
locals_max_string: int = LOCALS_MAX_STRING,
suppress: Iterable[Union[str, ModuleType]] = (),
):
if trace is None:
exc_type, exc_value, traceback = sys.exc_info()
Expand All @@ -233,6 +240,15 @@ def __init__(
self.locals_max_length = locals_max_length
self.locals_max_string = locals_max_string

self.suppress: Sequence[str] = []
for suppress_entity in suppress:
if not isinstance(suppress_entity, str):
path = os.path.dirname(suppress_entity.__file__)
else:
path = suppress_entity
path = os.path.normpath(os.path.abspath(path))
self.suppress.append(path)

@classmethod
def from_exception(
cls,
Expand All @@ -247,6 +263,7 @@ def from_exception(
indent_guides: bool = True,
locals_max_length: int = LOCALS_MAX_LENGTH,
locals_max_string: int = LOCALS_MAX_STRING,
suppress: Iterable[Union[str, ModuleType]] = (),
) -> "Traceback":
"""Create a traceback from exception info
Expand All @@ -263,6 +280,7 @@ 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.
Returns:
Traceback: A Traceback instance that may be printed.
Expand All @@ -280,6 +298,7 @@ def from_exception(
indent_guides=indent_guides,
locals_max_length=locals_max_length,
locals_max_string=locals_max_string,
suppress=suppress,
)

@classmethod
Expand Down Expand Up @@ -539,6 +558,9 @@ def render_locals(frame: Frame) -> Iterable[ConsoleRenderable]:
)

for first, frame in loop_first(stack.frames):
frame_filename = frame.filename
suppressed = any(frame_filename.startswith(path) for path in self.suppress)

text = Text.assemble(
path_highlighter(Text(frame.filename, style="pygments.string")),
(":", "pygments.text"),
Expand All @@ -553,6 +575,8 @@ 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)
Expand Down Expand Up @@ -592,6 +616,7 @@ def render_locals(frame: Frame) -> Iterable[ConsoleRenderable]:

if __name__ == "__main__": # pragma: no cover

import rich
from .console import Console

console = Console()
Expand Down Expand Up @@ -622,6 +647,6 @@ def error() -> None:
except:
slfkjsldkfj # type: ignore
except:
console.print_exception(show_locals=True)
console.print_exception(show_locals=True, suppress=[rich])

error()

0 comments on commit 1b49c9f

Please sign in to comment.