Isolate treescope renderings to improve performance and duplication-s…

…afety.

Includes a number of changes to improve the robustness and performance of treescope
output renderings:

- Isolates each rendering in the "Shadow DOM", ensuring CSS styles do not conflict
  between different renderings.
- Isolates JS definitions as properties of a custom container element, avoiding
  conflicts between different output versions.
- Adds duplication safety: making two copies of a treescope rendering will correctly
  render it twice instead of having interference between copies.
- Improves performance by adding CSS container content visibility annotations.

This change should significantly improve performance and responsiveness of Penzai's
visualizations when used in a notebook environment like JupyterLab, which does not
apply strong sandboxing between different output cells.

PiperOrigin-RevId: 649053637

penzai/treescope/__init__.py

@@ -39,7 +39,7 @@
 from . import foldable_representation
 from . import formatting_util
 from . import handlers
-from . import html_compression
+from . import html_encapsulation
 from . import html_escaping
 from . import renderer
 from . import repr_lib

penzai/treescope/arrayviz/arrayviz.py

@@ -30,7 +30,6 @@
 import json
 import os
 from typing import Any, Literal, Mapping, Sequence
-import uuid
 
 import jax
 import jax.numpy as jnp
@@ -66,6 +65,9 @@ def _html_setup() -> (
     set[part_interface.CSSStyleRule | part_interface.JavaScriptDefn]
 ):
   """Builds the setup HTML that should be included in any arrayviz output cell."""
+  arrayviz_src = html_escaping.heuristic_strip_javascript_comments(
+      load_arrayvis_javascript()
+  )
   return {
       part_interface.CSSStyleRule(html_escaping.without_repeated_whitespace("""
         .arrayviz_container {
@@ -121,9 +123,7 @@ def _html_setup() -> (
         }
       """)),
       part_interface.JavaScriptDefn(
-          html_escaping.heuristic_strip_javascript_comments(
-              load_arrayvis_javascript()
-          )
+          arrayviz_src + " this.getRootNode().host.defns.arrayviz = arrayviz;"
       ),
   }
 
@@ -215,9 +215,7 @@ def axis_spec_arg(i):
   for axis in slider_axes:
     sliced_axis_specs_arg.append(axis_spec_arg(axis))
 
-  fresh_id = "arrayviz" + uuid.uuid4().hex
   args_json = json.dumps({
-      "destinationId": fresh_id,
       "info": info,
       "arrayBase64": base64.b64encode(converted_array_data.tobytes()).decode(
           "ascii"
@@ -241,12 +239,27 @@ def axis_spec_arg(i):
       },
       "valueFormattingInstructions": formatting_instructions,
   })
+  # Note: We need to save the parent of the treescope-run-here element first,
+  # because it will be removed before the runSoon callback executes.
+  inner_fn = html_escaping.without_repeated_whitespace("""
+    const parent = this.parentNode;
+    const defns = this.getRootNode().host.defns;
+    defns.runSoon(() => {
+        const tpl = parent.querySelector('template.deferred_args');
+        const config = JSON.parse(
+            tpl.content.querySelector('script').textContent
+        );
+        tpl.remove();
+        defns.arrayviz.buildArrayvizFigure(parent, config);
+    });
+  """)
   src = (
-      f'<div id="{fresh_id}" class="arrayviz_container">'
+      '<div class="arrayviz_container">'
       '<span class="loading_message">Rendering array...</span>'
-      "</div>"
-      '<template class="treescope_run_soon"><script>'
-      f" arrayviz.buildArrayvizFigure({args_json})</script></template>"
+      f'<treescope-run-here><script type="application/octet-stream">{inner_fn}'
+      "</script></treescope-run-here>"
+      '<template class="deferred_args">'
+      f'<script type="application/json">{args_json}</script></template></div>'
   )
   return src
 
@@ -1471,19 +1484,24 @@ def integer_digitbox(
   if label_bottom is None:
     label_bottom = str(value)
 
-  fresh_id = "arrayviz" + uuid.uuid4().hex
   render_args = json.dumps({
       "value": value,
       "labelTop": label_top,
       "labelBottom": label_bottom,
-      "destinationId": fresh_id,
   })
   size_attr = html_escaping.escape_html_attribute(size)
+  # Note: We need to save the parent of the treescope-run-here element first,
+  # because it will be removed before the runSoon callback executes.
   src = (
-      f'<span id="{fresh_id}" class="inline_digitbox"'
-      f' style="font-size: {size_attr}"></span>'
-      '<template class="treescope_run_soon">'
-      f"<script>arrayviz.renderOneDigitbox({render_args});</script></template>"
+      f'<span class="inline_digitbox" style="font-size: {size_attr}">'
+      '<treescope-run-here><script type="application/octet-stream">'
+      "const parent = this.parentNode;"
+      "const defns = this.getRootNode().host.defns;"
+      "defns.runSoon(() => {"
+      f"defns.arrayviz.buildArrayvizFigure(parent, {render_args});"
+      "});"
+      "</script></treescope-run-here>"
+      "</span>"
   )
   return ArrayvizRendering(src)
 

penzai/treescope/arrayviz/js/arrayviz.js

@@ -17,10 +17,11 @@
 /**
  * @fileoverview Javascript implementation of arrayviz.
  * This file is inserted into the generated Treescope HTML when arrayviz is
- * used from Python. Compatible with the Closure JS compiler.
+ * used from Python.
  */
 
-window['arrayviz'] = (() => {
+/** @type {*} */
+const arrayviz = (() => {
   /**
    * Converts a base64 string to a Uint8 array.
    * @param {string} s
@@ -520,10 +521,8 @@ window['arrayviz'] = (() => {
   let AxisSpec;
 
   /* Delays an action until a destination element becomes visible. */
-  function _delayUntilVisible(destinationId, action) {
+  function _delayUntilVisible(destination, action) {
     // Trigger rendering as soon as the destination becomes visible.
-    const destination = /** @type {!Element} */
-        (document.getElementById(destinationId));
     const visiblityObserver = new IntersectionObserver((entries) => {
       if (entries[0].intersectionRatio > 0) {
         const loadingMarkers = destination.querySelectorAll('.loading_message');
@@ -541,7 +540,6 @@ window['arrayviz'] = (() => {
 
   /**
    * Renders an array to a destination object.
-   *  config.destinationId: ID of the HTML element to render into.
    *  config.info: Info for the figure; drawn on the bottom.
    *  config.arrayBase64: Base64-encoded array of either float32 or int32 data.
    *  config.arrayDtype: Either "float32" or "int32".
@@ -562,8 +560,8 @@ window['arrayviz'] = (() => {
    *    as continuous or discrete.
    *  config.valueFormattingInstructions: Instructions for rendering array
    *    indices and values on mouse hover/click.
+   * @param {!HTMLElement} destination The element to render into.
    * @param {{
-   *    destinationId: string,
    *    info: string,
    *    arrayBase64: string,
    *    arrayDtype: string,
@@ -576,13 +574,12 @@ window['arrayviz'] = (() => {
    *    valueFormattingInstructions: !Array<?>,
    * }} config Configuration for the setup.
    */
-  function buildArrayvizFigure(config) {
-    _delayUntilVisible(config.destinationId, () => {
-      _buildArrayvizFigure(config);
+  function buildArrayvizFigure(destination, config) {
+    _delayUntilVisible(destination, () => {
+      _buildArrayvizFigure(destination, config);
     });
   }
-  function _buildArrayvizFigure(config) {
-    const destinationId = config.destinationId;
+  function _buildArrayvizFigure(destination, config) {
     const info = config.info;
     const arrayBase64 = config.arrayBase64;
     const arrayDtype = config.arrayDtype;
@@ -884,9 +881,6 @@ window['arrayviz'] = (() => {
       return [array[dataIndex], valid_mask[dataIndex], result, row, col];
     }
 
-    // Set up the overall HTML structure for the rendering.
-    const destination = document.getElementById(destinationId);
-
     // HTML structure: The main figure, wrapped in a series of containers to
     // help with axis label layout.
     const container = /** @type {!HTMLDivElement} */ (
@@ -1015,7 +1009,7 @@ window['arrayviz'] = (() => {
     const datalist =
         /** @type {!HTMLDataListElement} */ (
             destination.appendChild(document.createElement('datalist')));
-    datalist.id = `${destinationId}-markers`;
+    datalist.id = "arrayviz-markers";
     for (const val of [1, 7, 14, 21]) {
       const datalistOption = /** @type {!HTMLOptionElement} */ (
           datalist.appendChild(document.createElement('option')));
@@ -1031,7 +1025,6 @@ window['arrayviz'] = (() => {
     zoomslider.setAttribute('list', datalist.id);
     const infodiv = /** @type {!HTMLDivElement} */ (
         destination.appendChild(document.createElement('div')));
-    infodiv.id = `${destinationId}-info`;
     infodiv.classList.add('info');
     infodiv.style.whiteSpace = 'pre';
     infodiv.append('Zoom: -');
@@ -1218,17 +1211,15 @@ window['arrayviz'] = (() => {
    *  config.value: Integer value to render.
    *  config.labelTop: Label to draw above the box.
    *  config.labelBottom: Label to draw below the box.
-   *  config.destinationId: ID of the element to render into.
+   * @param {!HTMLElement} destination The element to render into.
    * @param {{
    *    value: number,
    *    labelTop: string,
    *    labelBottom: string,
-   *    destinationId: string,
    * }} config Configuration for the digitbox.
    */
-  function renderOneDigitbox(config) {
-    _delayUntilVisible(config.destinationId, () => {
-      const destination = document.getElementById(config.destinationId);
+  function renderOneDigitbox(destination, config) {
+    _delayUntilVisible(destination, () => {
       const value = config.value;
       const labelTop = config.labelTop;
       const labelBottom = config.labelBottom;

penzai/treescope/default_renderer.py

@@ -254,6 +254,7 @@ def render_to_html(
     value: Any,
     roundtrip_mode: bool = False,
     ignore_exceptions: bool = False,
+    compressed: bool = True,
 ) -> str:
   """Renders an object to HTML using the default renderer.
 
@@ -263,11 +264,14 @@ def render_to_html(
     ignore_exceptions: Whether to catch errors during rendering of subtrees and
       show a fallback for those subtrees, instead of failing the entire
       renderer.
+    compressed: Whether to compress the output HTML.
 
   Returns:
     HTML source code for the foldable representation of the object.
   """
   foldable_ir = basic_parts.build_full_line_with_annotations(
       build_foldable_representation(value, ignore_exceptions=ignore_exceptions)
   )
-  return foldable_impl.render_to_html_as_root(foldable_ir, roundtrip_mode)
+  return foldable_impl.render_to_html_as_root(
+      foldable_ir, roundtrip_mode, compressed=compressed
+  )

penzai/treescope/figures.py

@@ -27,7 +27,6 @@
 from typing import Any
 
 from penzai.treescope import default_renderer
-from penzai.treescope import html_compression
 from penzai.treescope import html_escaping
 from penzai.treescope.foldable_representation import basic_parts
 from penzai.treescope.foldable_representation import embedded_iframe
@@ -45,10 +44,7 @@ class RendersAsRootInIPython(part_interface.RenderableTreePart):
 
   def _repr_html_(self) -> str:
     """Returns a rich HTML representation of this part."""
-    return html_compression.compress_html(
-        foldable_impl.render_to_html_as_root(self),
-        loading_message="(Loading...)",
-    )
+    return foldable_impl.render_to_html_as_root(self, compressed=True)
 
   def _repr_pretty_(self, p, cycle):
     """Builds a representation of this part for the IPython text prettyprinter."""

penzai/treescope/foldable_representation/basic_parts.py

@@ -1036,12 +1036,15 @@ def html_setup_parts(
       self, context: HtmlContextForSetup
   ) -> set[CSSStyleRule | JavaScriptDefn]:
     rule = html_escaping.without_repeated_whitespace(f"""
+        .indented_children {{
+            contain: content;
+        }}
         .indented_child:not({context.collapsed_selector} *)
         {{
             display: block;
             margin-left: calc(2ch - 1px);
         }}
-        .indented_children:not({context.collapsed_selector} *)
+        .stacked_children:not({context.collapsed_selector} *)
         {{
             display: block;
             border-left: dotted 1px #e0e0e0;

penzai/treescope/foldable_representation/common_styles.py

@@ -172,6 +172,9 @@ def html_setup_parts(
       self, context: HtmlContextForSetup
   ) -> set[CSSStyleRule | JavaScriptDefn]:
     rule = html_escaping.without_repeated_whitespace(f"""
+        .stacked_children {{
+          contain: content;
+        }}
         .stacked_children.colored_border
           > .stacked_child:not({context.collapsed_selector} *)
         {{
@@ -284,13 +287,17 @@ def _common_block_rules(
             {{
               padding-bottom: 0.15em;
               line-height: 1.5em;
+              z-index: 1;
+              position: relative;
             }}
           """)),
       "bottomline": CSSStyleRule(html_escaping.without_repeated_whitespace(f"""
             .bottomline:not({context.collapsed_selector} *)
             {{
               padding-top: 0.15em;
               line-height: 1.5em;
+              z-index: 1;
+              position: relative;
             }}
           """)),
       "hch_space_left": CSSStyleRule(

penzai/treescope/foldable_representation/embedded_iframe.py

@@ -120,7 +120,7 @@ def html_setup_parts(
         # itself to match its content.
         # But start with a minimum width of 80 characters.
         JavaScriptDefn(html_escaping.without_repeated_whitespace("""
-        window.treescope.resize_iframe_by_content = ((iframe) => {
+        this.getRootNode().host.defns.resize_iframe_by_content = ((iframe) => {
           iframe.height = 0;
           iframe.style.width = "80ch";
           iframe.style.overflow = "hidden";
@@ -165,5 +165,6 @@ def render_to_html(
     )
     stream.write(
         f'<div class="embedded_html"><iframe srcdoc="{srcdoc}"'
-        ' onload="treescope.resize_iframe_by_content(this)"></iframe></div>'
+        ' onload="this.getRootNode().host.defns.resize_iframe_by_content(this)">'
+        '</iframe></div>'
     )