Make exception printing customizable, add rich printer

hynek · Aug 28, 2021 · ad4e43c · ad4e43c
1 parent 1df431c
commit ad4e43c
Show file tree

Hide file tree

Showing 10 changed files with 190 additions and 41 deletions.
diff --git a/CHANGELOG.rst b/CHANGELOG.rst
@@ -29,8 +29,10 @@ Changes:
 ^^^^^^^^
 
 - ``structlog`` is now importable if ``sys.stdout`` is ``None`` (e.g. when running using ``pythonw``).
-- If the `better-exceptions <https://github.com/qix-/better-exceptions>`_ package is present, ``structlog.dev.ConsoleRenderer`` will now pretty-print exceptions using it.
-  Pass ``pretty_exceptions=False`` to disable.
+- Exception rendering in ``structlog.dev.ConsoleLogger`` is now configurable using the ``exception_formatter`` setting.
+  If either the `rich <https://github.com/willmcgugan/rich>`_ or the `better-exceptions <https://github.com/qix-/better-exceptions>`_ package is present, ``structlog`` will use them for pretty-printing tracebacks.
+  ``rich`` takes precedence over ``better-exceptions`` if both are present.
+
   This only works if ``format_exc_info`` is **absent** in the processor chain.
 - ``structlog.threadlocal.get_threadlocal()`` and ``structlog.contextvars.get_threadlocal()`` can now be used to get a copy of the current thread-local/context-local context that has been bound using ``structlog.threadlocal.bind_threadlocal()`` and ``structlog.contextvars.bind_contextvars()``.
 - ``structlog.threadlocal.get_merged_threadlocal(bl)`` and ``structlog.contextvars.get_merged_contextvars(bl)`` do the same, but also merge the context from a bound logger *bl*.

diff --git a/docs/_static/console_renderer.png b/docs/_static/console_renderer.png
diff --git a/docs/api.rst b/docs/api.rst
@@ -73,6 +73,10 @@ API Reference
 .. autoclass:: ConsoleRenderer
    :members: get_default_level_styles
 
+.. autofunction:: plain_traceback
+.. autofunction:: rich_traceback
+.. autofunction:: better_traceback
+
 .. autofunction:: set_exc_info
 
 

diff --git a/docs/conf.py b/docs/conf.py
@@ -72,6 +72,7 @@ def find_version(*file_paths):
     ("py:class", "PlainFileObserver"),
     ("py:class", "TLLogger"),
     ("py:class", "TextIO"),
+    ("py:class", "traceback"),
     ("py:class", "structlog._base.BoundLoggerBase"),
     ("py:class", "structlog.dev._Styles"),
     ("py:class", "structlog.types.EventDict"),

diff --git a/docs/development.rst b/docs/development.rst
@@ -4,7 +4,10 @@ Development
 To make development a more pleasurable experience, ``structlog`` comes with the `structlog.dev` module.
 
 The highlight is `structlog.dev.ConsoleRenderer` that offers nicely aligned and colorful (requires the `colorama package <https://pypi.org/project/colorama/>`_ if on Windows) console output.
-If the `better-exceptions <https://github.com/Qix-/better-exceptions>`_ package is installed, it will also pretty-print exceptions with helpful contextual data.
+
+If one of the `rich <https://rich.readthedocs.io/>`_ or  `better-exceptions <https://github.com/Qix-/better-exceptions>`_ packages is installed, it will also pretty-print exceptions with helpful contextual data. ``rich`` takes precedence over ``better-exceptions``, but you can configure it by passing `structlog.dev.plain_traceback` or `structlog.dev.better_traceback` for the ``exception_formatter`` parameter of `ConsoleRenderer`.
+
+The following output is rendered using ``rich``:
 
 .. figure:: _static/console_renderer.png
    :alt: Colorful console output by ConsoleRenderer.

diff --git a/docs/getting-started.rst b/docs/getting-started.rst
@@ -10,9 +10,9 @@ Installation
 
    $ pip install structlog
 
-If you'd like colorful output and pretty exceptions in development (you know you do!), install using::
+If you want pretty exceptions in development (you know you do!), additionally install either `rich <https://github.com/willmcgugan/rich>`_ or `better-exceptions <https://github.com/qix-/better-exceptions>`_. Try both to find out which one you like better -- the screenshot in the README and docs homepage is rendered by ``rich``.
 
-   $ pip install structlog colorama better-exceptions
+On Windows, you also have to install `colorama`_ if you want colorful output beside exceptions.
 
 
 Your First Log Entry

diff --git a/src/structlog/dev.py b/src/structlog/dev.py
@@ -12,10 +12,16 @@
 import warnings
 
 from io import StringIO
-from typing import Any, Optional, Type, Union
+from typing import Any, Optional, TextIO, Type, Union
 
 from ._frames import _format_exception
-from .types import EventDict, Protocol, WrappedLogger
+from .types import (
+    EventDict,
+    ExceptionFormatter,
+    ExcInfo,
+    Protocol,
+    WrappedLogger,
+)
 
 
 try:
@@ -28,8 +34,21 @@
 except ImportError:
     better_exceptions = None
 
+try:
+    import rich
+
+    from rich.console import Console
+    from rich.traceback import Traceback
+except ImportError:
+    rich = None  # type: ignore
+
 
-__all__ = ["ConsoleRenderer"]
+__all__ = [
+    "ConsoleRenderer",
+    "plain_traceback",
+    "rich_traceback",
+    "better_traceback",
+]
 
 _IS_WINDOWS = sys.platform == "win32"
 
@@ -137,13 +156,64 @@ class _PlainStyles:
     kv_value = ""
 
 
+def plain_traceback(sio: TextIO, exc_info: ExcInfo) -> None:
+    """
+    "Pretty"-print *exc_info* to *sio* using our own plain formatter.
+
+    To be passed into `ConsoleRenderer`'s ``exception_formatter`` argument.
+
+    Used by default if neither ``rich`` not ``better-exceptions`` are present.
+
+    .. versionadded:: 21.2
+    """
+    sio.write("\n" + _format_exception(exc_info))
+
+
+def rich_traceback(sio: TextIO, exc_info: ExcInfo) -> None:
+    """
+    Pretty-print *exc_info* to *sio* using the ``rich`` package.
+
+    To be passed into `ConsoleRenderer`'s ``exception_formatter`` argument.
+
+    Used by default if ``rich`` is installed.
+
+    .. versionadded:: 21.2
+    """
+    sio.write("\n")
+    Console(file=sio, color_system="truecolor").print(
+        Traceback.from_exception(*exc_info, show_locals=True)
+    )
+
+
+def better_traceback(sio: TextIO, exc_info: ExcInfo) -> None:
+    """
+    Pretty-print *exc_info* to *sio* using the ``better-exceptions`` package.
+
+    To be passed into `ConsoleRenderer`'s ``exception_formatter`` argument.
+
+    Used by default if ``better-exceptions`` is installed and ``rich`` is
+    absent.
+
+    .. versionadded:: 21.2
+    """
+    sio.write("\n" + "".join(better_exceptions.format_exception(*exc_info)))
+
+
+if rich is not None:
+    default_exception_formatter = rich_traceback
+elif better_exceptions is not None:  # type: ignore
+    default_exception_formatter = better_traceback
+else:
+    default_exception_formatter = plain_traceback
+
+
 class ConsoleRenderer:
     """
     Render ``event_dict`` nicely aligned, possibly in colors, and ordered.
 
     If ``event_dict`` contains a true-ish ``exc_info`` key, it will be
-    rendered *after* the log line. If better-exceptions_ is present, in
-    colors and with extra context.
+    rendered *after* the log line. If rich_ or better-exceptions_ are present,
+    in colors and with extra context.
 
     :param pad_event: Pad the event to this many characters.
     :param colors: Use colors for a nicer output. `True` by default if
@@ -160,14 +230,17 @@ class ConsoleRenderer:
         must be a dict from level names (strings) to colorama styles. The
         default can be obtained by calling
         `ConsoleRenderer.get_default_level_styles`
-    :param pretty_exceptions: Render exceptions with colors and extra
-        information. `True` by default if better-exceptions_ is installed.
+    :param exception_formatter: A callable to render ``exc_infos``. If rich_
+        or better-exceptions_ are installed, they are used for pretty-printing
+        by default (rich_ taking precendence). You can also manually set it to
+        `plain_traceback`, `better_traceback`, `rich_traceback`, or implement
+        your own.
 
-    Requires the colorama_ package if *colors* is `True`, and the
-    better-exceptions_ package if *pretty_exceptions* is `True`.
+    Requires the colorama_ package if *colors* is `True` **on Windows**.
 
     .. _colorama: https://pypi.org/project/colorama/
     .. _better-exceptions: https://pypi.org/project/better-exceptions/
+    .. _rich: https://pypi.org/project/rich/
 
     .. versionadded:: 16.0
     .. versionadded:: 16.1 *colors*
@@ -182,13 +255,14 @@ class ConsoleRenderer:
        anymore because it breaks rendering.
     .. versionchanged:: 21.1 It is additionally possible to set the logger name
        using the ``logger_name`` key in the ``event_dict``.
-    .. versionadded:: 21.2 *pretty_exceptions*
+    .. versionadded:: 21.2 *exception_formatter*
     .. versionchanged:: 21.2 `ConsoleRenderer` now handles the ``exc_info``
        event dict key itself. Do **not** use the
        `structlog.processors.format_exc_info` processor together with
        `ConsoleRenderer` anymore! It will keep working, but you can't have
-       pretty exceptions and a warning will be raised if you ask for them.
-    .. versionchanged:: 21.3 The colors keyword now defaults to True on
+       customize exception formatting and a warning will be raised if you ask
+       for it.
+    .. versionchanged:: 21.2 The colors keyword now defaults to True on
        non-Windows systems, and either True or False in Windows depending on
        whether colorama is installed.
     """
@@ -200,7 +274,7 @@ def __init__(
         force_colors: bool = False,
         repr_native_str: bool = False,
         level_styles: Optional[Styles] = None,
-        pretty_exceptions: bool = (better_exceptions is not None),
+        exception_formatter: ExceptionFormatter = default_exception_formatter,
     ):
         styles: Styles
         if colors:
@@ -241,7 +315,7 @@ def __init__(
         )
 
         self._repr_native_str = repr_native_str
-        self._pretty_exceptions = pretty_exceptions
+        self._exception_formatter = exception_formatter
 
     def _repr(self, val: Any) -> str:
         """
@@ -331,17 +405,11 @@ def __call__(
             if not isinstance(exc_info, tuple):
                 exc_info = sys.exc_info()
 
-            if self._pretty_exceptions:
-                sio.write(
-                    "\n"
-                    + "".join(better_exceptions.format_exception(*exc_info))
-                )
-            else:
-                sio.write("\n" + _format_exception(exc_info))
+            self._exception_formatter(sio, exc_info)
         elif exc is not None:
-            if self._pretty_exceptions:
+            if self._exception_formatter is not plain_traceback:
                 warnings.warn(
-                    "Remove `render_exc_info` from your processor chain "
+                    "Remove `format_exc_info` from your processor chain "
                     "if you want pretty exceptions."
                 )
             sio.write("\n" + exc)

diff --git a/src/structlog/types.py b/src/structlog/types.py
@@ -17,6 +17,7 @@
     Mapping,
     MutableMapping,
     Optional,
+    TextIO,
     Tuple,
     Type,
     Union,
@@ -81,6 +82,12 @@
 """
 
 
+ExceptionFormatter = Callable[[TextIO, ExcInfo], None]
+"""
+An callable that pretty-prints an `ExcInfo` into a `StringIO.
+"""
+
+
 @runtime_checkable
 class BindableLogger(Protocol):
     """

diff --git a/tests/test_dev.py b/tests/test_dev.py
@@ -5,6 +5,8 @@
 import pickle
 import sys
 
+from io import StringIO
+
 import pytest
 
 from structlog import dev
@@ -26,7 +28,9 @@ def test_negative(self):
 
 @pytest.fixture(name="cr")
 def _cr():
-    return dev.ConsoleRenderer(colors=dev._use_colors)
+    return dev.ConsoleRenderer(
+        colors=dev._use_colors, exception_formatter=dev.plain_traceback
+    )
 
 
 @pytest.fixture(name="styles")
@@ -204,23 +208,27 @@ def test_key_values(self, cr, styles, padded):
             + styles.reset
         ) == rv
 
-    def test_exception_rendered(self, cr, padded, recwarn):
+    @pytest.mark.parametrize("wrap", [True, False])
+    def test_exception_rendered(self, cr, padded, recwarn, wrap):
         """
         Exceptions are rendered after a new line if they are already rendered
         in the event dict.
 
-        A warning is emitted if pretty exceptions are active.
+        A warning is emitted if exception printing is "customized".
         """
         exc = "Traceback:\nFake traceback...\nFakeError: yolo"
 
+        # Wrap the formatter to provoke the warning.
+        if wrap:
+            cr._exception_formatter = lambda s, ei: dev.plain_tracebacks(s, ei)
         rv = cr(None, None, {"event": "test", "exception": exc})
 
         assert (padded + "\n" + exc) == rv
 
-        if cr._pretty_exceptions:
+        if wrap:
             (w,) = recwarn.list
             assert (
-                "Remove `render_exc_info` from your processor chain "
+                "Remove `format_exc_info` from your processor chain "
                 "if you want pretty exceptions.",
             ) == w.message.args
 
@@ -289,10 +297,7 @@ def test_everything(self, cr, styles, padded, explicit_ei):
                 rv = cr(None, None, ed)
                 ei = sys.exc_info()
 
-        if dev.better_exceptions:
-            exc = "".join(dev.better_exceptions.format_exception(*ei))
-        else:
-            exc = dev._format_exception(ei)
+        exc = dev._format_exception(ei)
 
         assert (
             styles.timestamp
@@ -431,3 +436,53 @@ def test_set_it(self):
         exception.
         """
         assert {"exc_info": True} == dev.set_exc_info(None, "exception", {})
+
+
+@pytest.mark.skipif(dev.rich is None, reason="Needs rich.")
+class TestRichTraceback:
+    def test_default(self):
+        """
+        If rich is present, it's the default.
+        """
+        assert dev.default_exception_formatter is dev.rich_traceback
+
+    def test_does_not_blow_up(self):
+        """
+        We trust rich to do the right thing, so we just exercise the function
+        and check the first new line that we add manually is present.
+        """
+        sio = StringIO()
+        try:
+            0 / 0
+        except ZeroDivisionError:
+            dev.rich_traceback(sio, sys.exc_info())
+
+        assert sio.getvalue().startswith("\n")
+
+
+@pytest.mark.skipif(
+    dev.better_exceptions is None, reason="Needs better-exceptions."
+)
+class TestBetterTraceback:
+    def test_default(self):
+        """
+        If better-exceptions is present and rich is NOT present, it's the
+        default.
+        """
+        assert (
+            dev.rich is not None
+            or dev.default_exception_formatter is dev.better_traceback
+        )
+
+    def test_does_not_blow_up(self):
+        """
+        We trust better-exceptions to do the right thing, so we just exercise
+        the function.
+        """
+        sio = StringIO()
+        try:
+            0 / 0
+        except ZeroDivisionError:
+            dev.better_traceback(sio, sys.exc_info())
+
+        assert sio.getvalue().startswith("\n")