Merge pull request #303 from taliraj/handler-with-tracebacks

handle exceptions with rich's logging handler

docs/source/logging.rst

@@ -20,4 +20,17 @@ Rich logs won't render :ref:`console_markup` in logging by default as most libra
 
     log.error("[bold red blink]Server is shutting down![/]", extra={"markup": True})
 
+
+Handle exceptions
+-------------------
+
+Rich's :class:`~rich.logging.RichHandler` class can be configured to handle exceptions. Here's an example::
+
+    handler = RichHandler(handle_tracebacks=True)
+
+    try:
+        1 / 0
+    except ZeroDivisionError:
+        log.exception("Exception message")
+
 There are a number of options you can use to configure logging output, see the :class:`~rich.logging.RichHandler` reference for details.

rich/logging.py

@@ -9,6 +9,7 @@ from ._log_render import LogRender
 from .console import Console
 from .highlighter import Highlighter, ReprHighlighter
 from .text import Text
+from .traceback import Traceback
 
 
 class RichHandler(Handler):
@@ -29,6 +30,11 @@ class RichHandler(Handler):
         enable_link_path (bool, optional): Enable terminal link of path column to file. Defaults to True.
         highlighter (Highlighter, optional): Highlighter to style log messages, or None to use ReprHighlighter. Defaults to None.
         markup (bool, optional): Enable console markup in log messages. Defaults to False.
+        handle_tracebacks (bool, optional): Enable rich tracebacks with syntax highlighting and formatting. Defaults to False.
+        tracebacks_width (Optional[int], optional): Number of characters used to render tracebacks code. Defaults to 88.
+        tracebacks_extra_lines (int, optional): Additional lines of code to render tracebacks. Defaults to 3.
+        tracebacks_theme (str, optional): Override pygments theme used in traceback
+        tracebacks_word_wrap (bool, optional): Enable word wrapping of long tracebacks lines. Defaults to False.
 
     """
 
@@ -55,6 +61,11 @@ class RichHandler(Handler):
         enable_link_path: bool = True,
         highlighter: Highlighter = None,
         markup: bool = False,
+        handle_tracebacks: bool = False,
+        tracebacks_width: Optional[int] = 88,
+        tracebacks_extra_lines: int = 3,
+        tracebacks_theme: Optional[str] = None,
+        tracebacks_word_wrap: bool = False,
     ) -> None:
         super().__init__(level=level)
         self.console = console or get_console()
@@ -64,6 +75,11 @@ class RichHandler(Handler):
         )
         self.enable_link_path = enable_link_path
         self.markup = markup
+        self.handle_tracebacks = handle_tracebacks
+        self.tracebacks_width = tracebacks_width
+        self.tracebacks_extra_lines = tracebacks_extra_lines
+        self.tracebacks_theme = tracebacks_theme
+        self.tracebacks_word_wrap = tracebacks_word_wrap
 
     def emit(self, record: LogRecord) -> None:
         """Invoked by logging."""
@@ -76,6 +92,17 @@ class RichHandler(Handler):
         level = Text()
         level.append(record.levelname, log_style)
 
+        traceback = None
+        if self.handle_tracebacks and record.exc_info:
+            traceback = Traceback.from_exception(
+                *record.exc_info,
+                width=self.tracebacks_width,
+                extra_lines=self.tracebacks_extra_lines,
+                theme=self.tracebacks_theme,
+                word_wrap=self.tracebacks_word_wrap,
+            )
+            message = record.getMessage()
+
         use_markup = (
             getattr(record, "markup") if hasattr(record, "markup") else self.markup
         )
@@ -92,7 +119,7 @@ class RichHandler(Handler):
         self.console.print(
             self._log_render(
                 self.console,
-                [message_text],
+                [message_text] if not traceback else [message_text, traceback],
                 log_time=log_time,
                 time_format=time_format,
                 level=level,

tests/test_logging.py

@@ -1,5 +1,8 @@
 import io
+import sys
+import os
 import logging
+import pytest
 
 from rich.console import Console
 from rich.logging import RichHandler
@@ -16,6 +19,12 @@ logging.basicConfig(
 log = logging.getLogger("rich")
 
 
+skip_win = pytest.mark.skipif(
+    os.name == "nt",
+    reason="rendered differently on windows",
+)
+
+
 def make_log():
     log.debug("foo")
     render = handler.console.file.getvalue()
@@ -25,10 +34,61 @@ def make_log():
 def test_log():
     render = make_log()
     print(repr(render))
-    expected = "\x1b[2;36m[DATE]\x1b[0m\x1b[2;36m \x1b[0m\x1b[32mDEBUG\x1b[0m    foo                                           \x1b[2mtest_logging.py\x1b[0m\x1b[2m:20\x1b[0m\n"
+    expected = "\x1b[2;36m[DATE]\x1b[0m\x1b[2;36m \x1b[0m\x1b[32mDEBUG\x1b[0m    foo                                           \x1b[2mtest_logging.py\x1b[0m\x1b[2m:29\x1b[0m\n"
     assert render == expected
 
 
+@skip_win
+def test_exception():
+    console = Console(
+        file=io.StringIO(), force_terminal=True, width=80, color_system="truecolor"
+    )
+    handler_with_tracebacks = RichHandler(
+        console=console, enable_link_path=False, handle_tracebacks=True
+    )
+    log.addHandler(handler_with_tracebacks)
+
+    try:
+        1 / 0
+    except ZeroDivisionError:
+        log.exception("message")
+
+    render = handler_with_tracebacks.console.file.getvalue()
+    print(render)
+
+    assert render.count("\n") == 16
+    assert "ZeroDivisionError" in render
+    assert "message" in render
+    assert "division by zero" in render
+
+
+def test_exception_with_extra_lines():
+    console = Console(
+        file=io.StringIO(), force_terminal=True, width=80, color_system="truecolor"
+    )
+    handler_extra_lines = RichHandler(
+        console=console,
+        enable_link_path=False,
+        markup=True,
+        handle_tracebacks=True,
+        tracebacks_extra_lines=5,
+    )
+    log.addHandler(handler_extra_lines)
+
+    try:
+        1 / 0
+    except ZeroDivisionError:
+        log.exception("message")
+
+    render = handler_extra_lines.console.file.getvalue()
+    print(render)
+
+    assert render.count("\n") == 21
+    assert "ZeroDivisionError" in render
+    assert "message" in render
+    assert "division by zero" in render
+
+
 if __name__ == "__main__":
     render = make_log()
     print(render)