Improvements to figure rendering and sizing.

- Improves size detection logic to properly account for the size of
  Plotly figures.
- Adds an override to directly unwrap treescope figure objects without
  IFrame indirection.
- Fixes deferred rendering insertion to preserve the original DOM nodes
  and their event listeners, instead of using a copy without event listeners.
- Allows collapsing automatic visualizations when used with `replace=False`.

PiperOrigin-RevId: 670950158

treescope/_internal/handlers/autovisualizer_hook.py

@@ -21,6 +21,7 @@
 from treescope import lowering
 from treescope import renderers
 from treescope import rendering_parts
+from treescope._internal import figures_impl
 from treescope._internal import object_inspection
 from treescope._internal.api import autovisualize
 
@@ -61,7 +62,9 @@ def use_autovisualizer_if_present(
       ordinary_result = node_renderer(node, path)
 
     if isinstance(result, IPythonVisualization):
-      if isinstance(result.display_object, object_inspection.HasReprHtml):
+      if isinstance(result.display_object, figures_impl.TreescopeFigure):
+        ipy_rendering = result.display_object.treescope_part
+      elif isinstance(result.display_object, object_inspection.HasReprHtml):
         obj = result.display_object
 
         def _thunk(_):
@@ -126,12 +129,28 @@ def _thunk(_):
         )
       else:
         replace = False
-        rendering_and_annotations = rendering_parts.RenderableAndLineAnnotations(
-            renderable=rendering_parts.floating_annotation_with_separate_focus(
-                rendering_parts.in_outlined_box(ipy_rendering)
-            ),
-            annotations=rendering_parts.empty_part(),
+        internal_contents = (
+            rendering_parts.floating_annotation_with_separate_focus(
+                rendering_parts.fold_condition(
+                    collapsed=rendering_parts.comment_color(
+                        rendering_parts.text("(Visualization hidden)")
+                    ),
+                    expanded=ipy_rendering,
+                )
+            )
         )
+        rendering_and_annotations = (
+            rendering_parts.RenderableAndLineAnnotations(
+                renderable=rendering_parts.in_outlined_box(
+                    rendering_parts.build_custom_foldable_tree_node(
+                        contents=internal_contents,
+                        expand_state=rendering_parts.ExpandState.EXPANDED,
+                    ).renderable
+                ),
+                annotations=rendering_parts.empty_part(),
+            )
+        )
+
     else:
       assert isinstance(result, VisualizationFromTreescopePart)
       replace = True

treescope/_internal/parts/embedded_iframe.py

@@ -84,27 +84,36 @@ def html_setup_parts(
       self, context: HtmlContextForSetup
   ) -> set[CSSStyleRule | JavaScriptDefn]:
     rules = {
-        # Register an interaction observer to detect when the content of the
-        # iframe changes size, then update the size of the iframe element
-        # itself to match its content.
-        # But start with a minimum width of 80 characters.
+        # We want to make sure that the iframe is big enough to contain its
+        # content, but this depends on how the content is formatted. Currently
+        # we take the minimum of two strategies:
+        # - The computed size of the content when given a frame size of 80
+        #   characters wide (the initial iframe size).
+        # - The size of the content scrollbar when rendered at width 0 (to catch
+        #   elements that are overflowing and don't get counted by the computed
+        #   size)
+        # We also detect resizes of the internal content and update the size of
+        # the iframe accordingly.
         JavaScriptDefn(html_escaping.without_repeated_whitespace("""
         this.getRootNode().host.defns.resize_iframe_by_content = ((iframe) => {
-          iframe.height = 0;
-          iframe.style.width = "80ch";
-          iframe.style.overflow = "hidden";
-          iframe.contentDocument.scrollingElement.style.width = "fit-content";
-          iframe.contentDocument.scrollingElement.style.height = "fit-content";
-          iframe.contentDocument.scrollingElement.style.overflow = "hidden";
+          const scroller = iframe.contentDocument.scrollingElement;
+          scroller.style.width = "0";
+          scroller.style.height = "0";
+          const minWidth = scroller.scrollWidth;
+          const minHeight = scroller.scrollHeight;
+          scroller.style.width = "fit-content";
+          scroller.style.height = "fit-content";
+          scroller.style.overflow = "hidden";
           const observer = new ResizeObserver((entries) => {
-            console.log("resize", entries);
             const [entry] = entries;
-            const computedStyle = getComputedStyle(
-                iframe.contentDocument.scrollingElement);
-            iframe.style.width = `calc(4ch + ${computedStyle['width']})`;
-            iframe.style.height = `calc(${computedStyle['height']})`;
+            const computedStyle = getComputedStyle(scroller);
+            iframe.style.width =
+              `calc(4ch + max(${computedStyle['width']}, ${minWidth}px))`;
+            iframe.style.height =
+              `calc(max(${computedStyle['height']}, ${minHeight}px))`;
           });
-          observer.observe(iframe.contentDocument.scrollingElement);
+          observer.observe(scroller);
+          console.log("registered scroller", scroller);
         });
         """)),
         CSSStyleRule(html_escaping.without_repeated_whitespace("""
@@ -134,6 +143,7 @@ def render_to_html(
     )
     stream.write(
         f'<div class="embedded_html"><iframe srcdoc="{srcdoc}"'
+        ' style="width: 80ch; height: 0; overflow: hidden"'
         ' onload="this.getRootNode().host.defns.resize_iframe_by_content(this)">'
         '</iframe></div>'
     )

treescope/lowering.py

@@ -382,20 +382,28 @@ def _render_one(
       stream.write("</span></div>")
 
     all_ids = [deferred.placeholder.replacement_id for deferred in deferreds]
+    # It's sometimes important to preserve node identity when inserting
+    # deferred objects, for instance if we've already registered event listeners
+    # on some nodes. However, editing the DOM in place can be slow because it
+    # requires re-rendering the tree on every edit. To avoid this, we swap out
+    # the tree with a clone, edit the original tree, then swap the original
+    # tree back in.
     inner_script = (
         f"const targetIds = {json.dumps(all_ids)};"
         + html_escaping.without_repeated_whitespace("""
         const docroot = this.getRootNode();
         const treeroot = docroot.querySelector(".treescope_root");
+        const treerootClone = treeroot.cloneNode(true);
+        treeroot.replaceWith(treerootClone);
         const fragment = document.createDocumentFragment();
-        const treerootClone = fragment.appendChild(treeroot.cloneNode(true));
+        fragment.appendChild(treeroot);
         for (let i = 0; i < targetIds.length; i++) {
             let target = fragment.getElementById(targetIds[i]);
             let sourceDiv = docroot.querySelector("#for_" + targetIds[i]);
             target.replaceWith(sourceDiv.firstElementChild);
             sourceDiv.remove();
         }
-        treeroot.replaceWith(treerootClone);
+        treerootClone.replaceWith(treeroot);
         """)
     )
     stream.write(