Add additional error fallbacks during rendering.

tests/fixtures/treescope_examples_fixture.py

@@ -27,6 +27,7 @@
 
 import jax
 import torch
+import treescope
 
 
 class MyTestEnum(enum.Enum):
@@ -122,6 +123,7 @@ def wrapped_function():
 
 
 class SomethingCallable:
+
   def __call__(self, value: int) -> int:
     return value + 1
 
@@ -175,6 +177,47 @@ def __repr__(self):
     return "Non-idiomatic\nmultiline\nobject"
 
 
+class ObjectWithCustomHandler:
+
+  def __treescope_repr__(self, path, subtree_renderer):
+    del subtree_renderer
+    return treescope.rendering_parts.text(
+        f"<ObjectWithCustomHandler custom rendering! Path: {repr(path)}>"
+    )
+
+
+class ObjectWithCustomHandlerThatThrows:
+
+  def __treescope_repr__(self, path, subtree_renderer):
+    del path, subtree_renderer
+    raise RuntimeError("Simulated treescope_repr failure!")
+
+  def __repr__(self):
+    return "<Fallback repr for ObjectWithCustomHandlerThatThrows>"
+
+
+class ObjectWithReprThatThrows:
+
+  def __repr__(self):
+    raise RuntimeError("Simulated repr failure!")
+
+
+class ObjectWithCustomHandlerThatThrowsDeferred:
+
+  def __treescope_repr__(self, path, subtree_renderer):
+    del path, subtree_renderer
+    def _internal_main_thunk(layout_decision):
+      del layout_decision
+      raise RuntimeError("Simulated deferred treescope_repr failure!")
+
+    return treescope.lowering.maybe_defer_rendering(
+        main_thunk=_internal_main_thunk,
+        placeholder_thunk=lambda: treescope.rendering_parts.text(
+            "<deferred placeholder>"
+        ),
+    )
+
+
 class SomePyTorchModule(torch.nn.Module):
   """A basic PyTorch module to test rendering."""
 

tests/renderer_test.py

@@ -395,17 +395,13 @@ def hook_that_crashes(node, path, node_renderer):
           testcase_name="well_known_function",
           target=treescope.render_to_text,
           expected_collapsed="render_to_text",
-          expected_roundtrip_collapsed=(
-              "treescope.render_to_text"
-          ),
+          expected_roundtrip_collapsed="treescope.render_to_text",
       ),
       dict(
           testcase_name="well_known_type",
           target=treescope.IPythonVisualization,
           expected_collapsed="IPythonVisualization",
-          expected_roundtrip_collapsed=(
-              "treescope.IPythonVisualization"
-          ),
+          expected_roundtrip_collapsed="treescope.IPythonVisualization",
       ),
       dict(
           testcase_name="ast_nodes",
@@ -426,6 +422,21 @@ def hook_that_crashes(node, path, node_renderer):
                 ),
               )"""),
       ),
+      dict(
+          testcase_name="custom_handler",
+          target=[fixture_lib.ObjectWithCustomHandler()],
+          expected_collapsed=(
+              "[<ObjectWithCustomHandler custom rendering! Path: '[0]'>]"
+          ),
+      ),
+      dict(
+          testcase_name="custom_handler_that_throws",
+          target=[fixture_lib.ObjectWithCustomHandlerThatThrows()],
+          ignore_exceptions=True,
+          expected_collapsed=(
+              "[<Fallback repr for ObjectWithCustomHandlerThatThrows>]"
+          ),
+      ),
       dict(
           testcase_name="dtype_standard",
           target=np.dtype(np.float32),
@@ -533,6 +544,7 @@ def test_object_rendering(
       expected_roundtrip: str | None = None,
       expected_roundtrip_collapsed: str | None = None,
       expand_depth: int = 1,
+      ignore_exceptions: bool = False,
   ):
     if target_builder is not None:
       assert target is None
@@ -541,7 +553,9 @@ def test_object_rendering(
     renderer = treescope.active_renderer.get()
     # Render it to IR.
     rendering = rendering_parts.build_full_line_with_annotations(
-        renderer.to_foldable_representation(target)
+        renderer.to_foldable_representation(
+            target, ignore_exceptions=ignore_exceptions
+        )
     )
 
     # Collapse all foldables.
@@ -671,6 +685,69 @@ def test_fallback_repr_one_line(self):
             ]"""),
     )
 
+  def test_fallback_repr_after_error(self):
+    target = [fixture_lib.ObjectWithCustomHandlerThatThrows()]
+    renderer = treescope.active_renderer.get()
+
+    with self.assertRaisesWithLiteralMatch(
+        RuntimeError, "Simulated treescope_repr failure!"
+    ):
+      renderer.to_foldable_representation(target)
+
+    rendering = rendering_parts.build_full_line_with_annotations(
+        renderer.to_foldable_representation(target, ignore_exceptions=True)
+    )
+
+    layout_algorithms.expand_to_depth(rendering, 0)
+    self.assertEqual(
+        lowering.render_to_text_as_root(rendering),
+        "[<Fallback repr for ObjectWithCustomHandlerThatThrows>]",
+    )
+    layout_algorithms.expand_to_depth(rendering, 2)
+    self.assertEqual(
+        lowering.render_to_text_as_root(rendering),
+        textwrap.dedent(f"""\
+            [
+              <Fallback repr for ObjectWithCustomHandlerThatThrows>,  # {object.__repr__(target[0])}
+            ]"""),
+    )
+
+  def test_ignore_exceptions_in_deferred(self):
+    target = [fixture_lib.ObjectWithCustomHandlerThatThrowsDeferred()]
+    renderer = treescope.active_renderer.get()
+
+    with self.assertRaisesWithLiteralMatch(
+        RuntimeError, "Simulated deferred treescope_repr failure!"
+    ):
+      renderer.to_foldable_representation(target)
+
+    with lowering.collecting_deferred_renderings() as deferreds:
+      foldable_ir = rendering_parts.build_full_line_with_annotations(
+          renderer.to_foldable_representation(target)
+      )
+
+    # It's difficult to test the IPython wrapper so we instead test the internal
+    # helper function that produces the streaming HTML output.
+    html_parts = lowering._render_to_html_as_root_streaming(
+        root_node=foldable_ir,
+        roundtrip=False,
+        deferreds=deferreds,
+        ignore_exceptions=True,
+    )
+    self.assertContainsInOrder(
+        [
+            "[",
+            "&lt;RuntimeError during deferred rendering",
+            "Traceback",
+            "in _internal_main_thunk",
+            "raise RuntimeError",
+            "RuntimeError: Simulated deferred treescope_repr failure!",
+            "&gt;",
+            "]",
+        ],
+        "".join(html_parts),
+    )
+
   def test_fallback_repr_multiline_idiomatic(self):
     target = [fixture_lib.UnknownObjectWithMultiLineRepr()]
     renderer = treescope.active_renderer.get()
@@ -738,6 +815,33 @@ def test_fallback_repr_basic(self):
             ]"""),
     )
 
+  def test_failsafe_for_throw_in_repr(self):
+    target = [fixture_lib.ObjectWithReprThatThrows()]
+    renderer = treescope.active_renderer.get()
+
+    with self.assertRaisesWithLiteralMatch(
+        RuntimeError, "Simulated repr failure!"
+    ):
+      renderer.to_foldable_representation(target)
+
+    rendering = rendering_parts.build_full_line_with_annotations(
+        renderer.to_foldable_representation(target, ignore_exceptions=True)
+    )
+
+    layout_algorithms.expand_to_depth(rendering, 0)
+    self.assertEqual(
+        lowering.render_to_text_as_root(rendering),
+        f"[{object.__repr__(target[0])}]",
+    )
+    layout_algorithms.expand_to_depth(rendering, 2)
+    self.assertEqual(
+        lowering.render_to_text_as_root(rendering),
+        textwrap.dedent(f"""\
+            [
+              {object.__repr__(target[0])},  # Error occured while formatting this object.
+            ]"""),
+    )
+
   def test_shared_values(self):
     shared = ["bar"]
     target = [shared, shared, {"foo": shared}]
@@ -803,9 +907,7 @@ def inner_autovisualizer(node, path):
             ),
         )
 
-    with treescope.active_autovisualizer.set_scoped(
-        autovisualizer_for_test
-    ):
+    with treescope.active_autovisualizer.set_scoped(autovisualizer_for_test):
       renderer = treescope.active_renderer.get()
       rendering = rendering_parts.build_full_line_with_annotations(
           renderer.to_foldable_representation(target)

treescope/_internal/api/ipython_integration.py

@@ -94,6 +94,7 @@ def _display_and_maybe_steal(
           roundtrip=roundtrip_mode,
           compressed=compress_html,
           stealable=stealable,
+          ignore_exceptions=ignore_exceptions,
       )
     else:
       rendering = lowering.render_to_html_as_root(

treescope/lowering.py

@@ -23,6 +23,7 @@
 import contextlib
 import io
 import json
+import traceback
 from typing import Any, Callable, Iterator, Sequence
 import uuid
 
@@ -162,8 +163,10 @@ def render_to_text_as_root(
     Text for the rendered node.
   """
   if strip_whitespace_lines and not strip_trailing_whitespace:
-    raise ValueError("strip_whitespace_lines must be False if "
-                     "strip_trailing_whitespace is False.")
+    raise ValueError(
+        "strip_whitespace_lines must be False if "
+        "strip_trailing_whitespace is False."
+    )
 
   stream = io.StringIO()
   root_node.render_to_text(
@@ -223,13 +226,16 @@ def _render_to_html_as_root_streaming(
     root_node: rendering_parts.RenderableTreePart,
     roundtrip: bool,
     deferreds: Sequence[foldable_impl.DeferredWithThunk],
+    ignore_exceptions: bool = False,
 ) -> Iterator[str]:
   """Helper function: renders a root node to HTML one step at a time.
 
   Args:
     root_node: The root node to render.
     roundtrip: Whether to render in roundtrip mode.
     deferreds: Sequence of deferred objects to render and splice in.
+    ignore_exceptions: Whether to ignore exceptions during deferred rendering,
+      replacing them with error markers.
 
   Yields:
     HTML source for the rendered node, followed by logic to substitute each
@@ -347,7 +353,26 @@ def _render_one(
           layout_decision = deferred.placeholder.child.get_expand_state()
         else:
           layout_decision = None
-        replacement_part = deferred.thunk(layout_decision)
+        try:
+          replacement_part = deferred.thunk(layout_decision)
+        except Exception as e:  # pylint: disable=broad-except
+          if not ignore_exceptions:
+            raise
+          exc_child = rendering_parts.fold_condition(
+              expanded=rendering_parts.indented_children(
+                  [rendering_parts.text(traceback.format_exc())]
+              ),
+          )
+          replacement_part = rendering_parts.error_color(
+              rendering_parts.build_custom_foldable_tree_node(
+                  label=rendering_parts.text(
+                      f"<{type(e).__name__} during deferred rendering"
+                  ),
+                  contents=rendering_parts.siblings(
+                      exc_child, rendering_parts.text(">")
+                  ),
+              ).renderable
+          )
       _render_one(
           replacement_part,
           deferred.placeholder.saved_at_beginning_of_line,
@@ -409,6 +434,7 @@ def display_streaming_as_root(
     roundtrip: bool = False,
     compressed: bool = True,
     stealable: bool = False,
+    ignore_exceptions: bool = False,
 ) -> str | None:
   """Displays a root node in an IPython notebook in a streaming fashion.
 
@@ -419,6 +445,8 @@ def display_streaming_as_root(
     compressed: Whether to compress the HTML for display.
     stealable: Whether to return an extra HTML snippet that allows the streaming
       rendering to be relocated after it is shown.
+    ignore_exceptions: Whether to ignore exceptions during deferred rendering,
+      replacing them with error markers.
 
   Returns:
     If ``stealable`` is True, a final HTML snippet which, if inserted into a
@@ -431,7 +459,7 @@ def display_streaming_as_root(
   import IPython.display  # pylint: disable=g-import-not-at-top
 
   render_iterator = _render_to_html_as_root_streaming(
-      root_node, roundtrip, deferreds
+      root_node, roundtrip, deferreds, ignore_exceptions=ignore_exceptions
   )
   encapsulated_iterator = html_encapsulation.encapsulate_streaming_html(
       render_iterator, compress=compressed, stealable=stealable

treescope/renderers.py

@@ -328,12 +328,18 @@ def _render_subtree(
               f"No handler registered for a node of type {type(node)}."
           )
           # Fall back to a basic `repr` so that we still render something even
-          # without a handler for it.
+          # without a handler for it. We use the object repr because a custom
+          # repr may still raise an exception if the object is in an invalid
+          # state.
           return rendering_parts.RenderableAndLineAnnotations(
-              renderable=rendering_parts.abbreviation_color(
-                  rendering_parts.text(repr(node))
+              renderable=rendering_parts.error_color(
+                  rendering_parts.text(object.__repr__(node))
+              ),
+              annotations=rendering_parts.comment_color(
+                  rendering_parts.text(
+                      "  # Error occured while formatting this object."
+                  )
               ),
-              annotations=rendering_parts.empty_part(),
           )
         else:
           raise ValueError(