Merge pull request #994 from nathanrpage97/feat/ipython-traceback

Add traceback feature for IPython

rich/live_render.py

@@ -1,11 +1,13 @@
-from typing import Literal, Optional, Tuple
+from typing import Optional, Tuple
 
+from typing_extensions import Literal
+
+from ._loop import loop_last
 from .console import Console, ConsoleOptions, RenderableType, RenderResult
 from .control import Control
 from .segment import Segment
-from .text import Text
 from .style import StyleType
-from ._loop import loop_last
+from .text import Text
 
 VerticalOverflowMethod = Literal["crop", "ellipsis", "visible"]
 
@@ -88,4 +90,4 @@ class LiveRender:
         for last, line in loop_last(lines):
             yield from line
             if not last:
-                yield _Segment.line()
+                yield _Segment.line()

rich/traceback.py

@@ -72,7 +72,7 @@ def install(
     def excepthook(
         type_: Type[BaseException],
         value: BaseException,
-        traceback: TracebackType,
+        traceback: Optional[TracebackType],
     ) -> None:
         traceback_console.print(
             Traceback.from_exception(
@@ -88,9 +88,55 @@ def install(
             )
         )
 
-    old_excepthook = sys.excepthook
-    sys.excepthook = excepthook
-    return old_excepthook
+    def ipy_excepthook_closure(ip) -> None:  # pragma: no cover
+        tb_data = {}  # store information about showtraceback call
+        default_showtraceback = ip.showtraceback  # keep reference of default traceback
+
+        def ipy_show_traceback(*args, **kwargs) -> None:
+            """wrap the default ip.showtraceback to store info for ip._showtraceback"""
+            nonlocal tb_data
+            tb_data = kwargs
+            default_showtraceback(*args, **kwargs)
+
+        def ipy_display_traceback(*args, is_syntax: bool = False, **kwargs) -> None:
+            """Internally called traceback from ip._showtraceback"""
+            nonlocal tb_data
+            exc_tuple = ip._get_exc_info()
+
+            # do not display trace on syntax error
+            tb: Optional[TracebackType] = None if is_syntax else exc_tuple[2]
+
+            # determine correct tb_offset
+            compiled = tb_data.get("running_compiled_code", False)
+            tb_offset = tb_data.get("tb_offset", 1 if compiled else 0)
+            # remove ipython internal frames from trace with tb_offset
+            for _ in range(tb_offset):
+                if tb is None:
+                    break
+                tb = tb.tb_next
+
+            excepthook(exc_tuple[0], exc_tuple[1], tb)
+            tb_data = {}  # clear data upon usage
+
+        # replace _showtraceback instead of showtraceback to allow ipython features such as debugging to work
+        # this is also what the ipython docs recommends to modify when subclassing InteractiveShell
+        ip._showtraceback = ipy_display_traceback
+        # add wrapper to capture tb_data
+        ip.showtraceback = ipy_show_traceback
+        ip.showsyntaxerror = lambda *args, **kwargs: ipy_display_traceback(
+            *args, is_syntax=True, **kwargs
+        )
+
+    try:  # pragma: no cover
+        # if wihin ipython, use customized traceback
+        ip = get_ipython()  # type: ignore
+        ipy_excepthook_closure(ip)
+        return sys.excepthook
+    except Exception:
+        # otherwise use default system hook
+        old_excepthook = sys.excepthook
+        sys.excepthook = excepthook
+        return old_excepthook
 
 
 @dataclass

tests/test_console.py

@@ -491,4 +491,4 @@ def test_no_nested_live():
     with pytest.raises(errors.LiveError):
         with console.status("foo"):
             with console.status("bar"):
-                pass
+                pass