Enable rendering of rich exception objects in traceback

CHANGELOG.md

@@ -5,6 +5,12 @@ 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).
 
+## Unreleased
+
+### Added
+
+- Exception objects that are `RichRenderable` are rendered, instead of using the class name and string representation https://github.com/Textualize/rich/pull/3325
+
 ## [13.7.1] - 2023-02-28
 
 ### Fixed

CONTRIBUTORS.md

@@ -17,6 +17,7 @@ The following people have contributed to the development of Rich:
 - [James Estevez](https://github.com/jstvz)
 - [Aryaz Eghbali](https://github.com/AryazE)
 - [Oleksis Fraga](https://github.com/oleksis)
+- [Pradyun Gedam](https://github.com/pradyunsg)
 - [Andy Gimblett](https://github.com/gimbo)
 - [Michał Górny](https://github.com/mgorny)
 - [Nok Lam Chan](https://github.com/noklam)

docs/source/traceback.rst

@@ -22,7 +22,7 @@ The :meth:`~rich.console.Console.print_exception` method will print a traceback
         console.print_exception(show_locals=True)
 
 The ``show_locals=True`` parameter causes Rich to display the value of local variables for each frame of the traceback.
- 
+
 See `exception.py <https://github.com/willmcgugan/rich/blob/master/examples/exception.py>`_ for a larger example.
 
 
@@ -63,7 +63,7 @@ Suppressing Frames
 
 If you are working with a framework (click, django etc), you may only be interested in seeing the code from your own application within the traceback. You can exclude framework code by setting the `suppress` argument on `Traceback`, `install`, `Console.print_exception`, and `RichHandler`, which should be a list of modules or str paths.
 
-Here's how you would exclude `click <https://click.palletsprojects.com/en/8.0.x/>`_ from Rich exceptions:: 
+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
@@ -96,3 +96,32 @@ Here's an example of printing a recursive error::
     except Exception:
         console.print_exception(max_frames=20)
 
+Rendering Rich Exceptions
+-------------------------
+
+You can create exceptions that implement :ref:`protocol`, which would be rendered when presented in a traceback.
+
+Here's an example that renders the exception's message in a :ref:`~rich.panel.Panel` with an ASCII box::
+
+    from rich.box import ASCII
+    from rich.console import Console
+    from rich.panel import Panel
+
+
+    class MyAwesomeException(Exception):
+        def __init__(self, message: str):
+            super().__init__(message)
+            self.message = message
+
+        def __rich__(self):
+            return Panel(self.message, title="My Awesome Exception", box=ASCII)
+
+
+    def do_something():
+        raise MyAwesomeException("Something went wrong")
+
+    console = Console()
+    try:
+        do_something()
+    except Exception:
+        console.print_exception()

rich/traceback.py

@@ -28,6 +28,7 @@
 
 from . import pretty
 from ._loop import loop_last
+from .abc import RichRenderable
 from .columns import Columns
 from .console import Console, ConsoleOptions, ConsoleRenderable, RenderResult, group
 from .constrain import Constrain
@@ -194,6 +195,7 @@ class _SyntaxError:
 class Stack:
     exc_type: str
     exc_value: str
+    exc_renderable: Optional[RichRenderable] = None
     syntax_error: Optional[_SyntaxError] = None
     is_cause: bool = False
     frames: List[Frame] = field(default_factory=list)
@@ -416,6 +418,8 @@ def safe_str(_object: Any) -> str:
                     line=exc_value.text or "",
                     msg=exc_value.msg,
                 )
+            if isinstance(exc_value, RichRenderable):
+                stack.exc_renderable = exc_value
 
             stacks.append(stack)
             append = stack.frames.append
@@ -544,6 +548,19 @@ def __rich_console__(
                     (f"{stack.exc_type}: ", "traceback.exc_type"),
                     highlighter(stack.syntax_error.msg),
                 )
+            elif stack.exc_renderable:
+                exc_renderable = Constrain(stack.exc_renderable, self.width)
+                with console.use_theme(traceback_theme):
+                    try:
+                        segments = tuple(console.render(renderable=exc_renderable))
+                    except Exception:
+                        yield Text("<exception rich render failed>")
+                        yield Text.assemble(
+                            (f"{stack.exc_type}: ", "traceback.exc_type"),
+                            highlighter(stack.exc_value),
+                        )
+                    else:
+                        yield from segments
             elif stack.exc_value:
                 yield Text.assemble(
                     (f"{stack.exc_type}: ", "traceback.exc_type"),
@@ -722,7 +739,11 @@ def render_locals(frame: Frame) -> Iterable[ConsoleRenderable]:
     from .console import Console
 
     console = Console()
-    import sys
+
+    class RichError(Exception):
+        def __rich_console__(self, console, options):
+            yield "[bold][red]ERROR[/]:[/] This is a [i]renderable[/] exception!"
+            yield Panel(self.args[0], expand=False, border_style="blue")
 
     def bar(a: Any) -> None:  # 这是对亚洲语言支持的测试。面对模棱两可的想法，拒绝猜测的诱惑
         one = 1
@@ -746,7 +767,7 @@ def error() -> None:
             try:
                 foo(0)
             except:
-                slfkjsldkfj  # type: ignore[name-defined]
+                raise RichError("Woah! Look at this text!")
         except:
             console.print_exception(show_locals=True)
 

tests/test_traceback.py

@@ -248,6 +248,27 @@ def test_traceback_console_theme_applies():
     assert f"\\x1b[38;2;{r};{g};{b}mTraceback \\x1b[0m" in repr(result)
 
 
+def test_rich_exception():
+    class RichError(Exception):
+        def __rich_console__(self, console, options):
+            yield f"[bold][red]error[/red]:[/bold] [red]this is red[/red]"
+
+    console = Console(
+        width=100,
+        file=io.StringIO(),
+        color_system="truecolor",
+        legacy_windows=False,
+    )
+    try:
+        raise RichError()
+    except Exception:
+        console.print_exception()
+    result = console.file.getvalue()
+    print(result)
+    assert "\x1b[1;31merror\x1b[0m\x1b[1m:\x1b[0m" in result
+    assert "\x1b[31mthis is red\x1b[0m" in result
+
+
 def test_broken_str():
     class BrokenStr(Exception):
         def __str__(self):
@@ -263,6 +284,36 @@ def __str__(self):
     assert "<exception str() failed>" in result
 
 
+def test_broken_rich_exception():
+    class BrokenRichError(Exception):
+        def __rich_console__(self, console, options):
+            raise Exception("broken")
+
+    console = Console(width=100, file=io.StringIO())
+    try:
+        raise BrokenRichError()
+    except Exception:
+        console.print_exception()
+    result = console.file.getvalue()
+    print(result)
+    assert "<exception rich render failed>" in result
+
+
+def test_broken_rich_bad_markup():
+    class BrokenRichError(Exception):
+        def __rich_console__(self, console, options):
+            yield "[red]broken[/green]"
+
+    console = Console(width=100, file=io.StringIO())
+    try:
+        raise BrokenRichError()
+    except Exception:
+        console.print_exception()
+    result = console.file.getvalue()
+    print(result)
+    assert "<exception rich render failed>" in result
+
+
 def test_guess_lexer():
     assert Traceback._guess_lexer("foo.py", "code") == "python"
     code_python = "#! usr/bin/env python\nimport this"