diff --git a/penzai/treescope/__init__.py b/penzai/treescope/__init__.py index 62a4806..b5965d4 100644 --- a/penzai/treescope/__init__.py +++ b/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 diff --git a/penzai/treescope/arrayviz/arrayviz.py b/penzai/treescope/arrayviz/arrayviz.py index 93fcc10..7340c81 100644 --- a/penzai/treescope/arrayviz/arrayviz.py +++ b/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'
' + '
' 'Rendering array...' - "
" - '" + f'" + '
' ) 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'' - '" + f'' + '" + "" ) return ArrayvizRendering(src) diff --git a/penzai/treescope/arrayviz/js/arrayviz.js b/penzai/treescope/arrayviz/js/arrayviz.js index ba028af..f0de037 100644 --- a/penzai/treescope/arrayviz/js/arrayviz.js +++ b/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; diff --git a/penzai/treescope/default_renderer.py b/penzai/treescope/default_renderer.py index 94cc820..0658591 100644 --- a/penzai/treescope/default_renderer.py +++ b/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,6 +264,7 @@ 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. @@ -270,4 +272,6 @@ def render_to_html( 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 + ) diff --git a/penzai/treescope/figures.py b/penzai/treescope/figures.py index 7a2e491..1c27447 100644 --- a/penzai/treescope/figures.py +++ b/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.""" diff --git a/penzai/treescope/foldable_representation/basic_parts.py b/penzai/treescope/foldable_representation/basic_parts.py index 2c8b1c4..ada6d88 100644 --- a/penzai/treescope/foldable_representation/basic_parts.py +++ b/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; diff --git a/penzai/treescope/foldable_representation/common_styles.py b/penzai/treescope/foldable_representation/common_styles.py index 24748c0..95a7ee9 100644 --- a/penzai/treescope/foldable_representation/common_styles.py +++ b/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,6 +287,8 @@ 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""" @@ -291,6 +296,8 @@ def _common_block_rules( {{ padding-top: 0.15em; line-height: 1.5em; + z-index: 1; + position: relative; }} """)), "hch_space_left": CSSStyleRule( diff --git a/penzai/treescope/foldable_representation/embedded_iframe.py b/penzai/treescope/foldable_representation/embedded_iframe.py index 976ba8b..429f583 100644 --- a/penzai/treescope/foldable_representation/embedded_iframe.py +++ b/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'
' + ' onload="this.getRootNode().host.defns.resize_iframe_by_content(this)">' + '' ) diff --git a/penzai/treescope/foldable_representation/foldable_impl.py b/penzai/treescope/foldable_representation/foldable_impl.py index c09874a..9eeecd4 100644 --- a/penzai/treescope/foldable_representation/foldable_impl.py +++ b/penzai/treescope/foldable_representation/foldable_impl.py @@ -28,7 +28,7 @@ import uuid from penzai.core import context -from penzai.treescope import html_compression +from penzai.treescope import html_encapsulation from penzai.treescope import html_escaping from penzai.treescope.foldable_representation import basic_parts from penzai.treescope.foldable_representation import common_styles @@ -320,9 +320,10 @@ def html_setup_parts( return possible ? possible : target; }; - window.treescope.expand_and_scroll_to = ( + const defns = this.getRootNode().host.defns; + defns.expand_and_scroll_to = ( (linkelement, target_path) => { - const root = window.treescope.get_treescope_root(linkelement); + const root = linkelement.getRootNode().shadowRoot; const target = _get_target(root, target_path); /* Expand all of its parents. */ let may_need_expand = target.parentElement; @@ -350,9 +351,9 @@ def html_setup_parts( } ); - window.treescope.handle_hyperlink_mouse = ( + defns.handle_hyperlink_mouse = ( (linkelement, event, target_path) => { - const root = window.treescope.get_treescope_root(linkelement); + const root = linkelement.getRootNode().shadowRoot; const target = _get_target(root, target_path); if (event.type == "mouseover") { target.classList.add("hyperlink_remote_hover"); @@ -412,12 +413,12 @@ def render_to_html( stream.write( html_escaping.without_repeated_whitespace( '' ) ) @@ -480,7 +481,7 @@ def html_setup_parts( font_data_url = "data:font/woff2;base64,d09GMgABAAAAAAIQAAoAAAAABRwAAAHFAAEAAAAAAAAAAAAAAAAAAAAAAAAAAAAABmAAgkoKdIEBCwYAATYCJAMIBCAFgxIHLhtsBMieg3FDX2LI6YvSpuiPM5T1JXIRVMvWMztPyFFC+WgkTiBD0hiDQuJEdGj0Hb/fvIdpqK6hqWiiQuMnGhHfUtAU6BNr4AFInkE6cUuun+R5qcskwvfFl/qxgEo8gbJwG81HA/nAR5LrrJ1R+gz0Rd0AJf1gN7CwGj2g0oyuR77mE16wHX9OggpeTky4eIbz5cbrOGtaAgQINwDasysQuIIXWEFwAPQpIYdU//+g7T7X3t0fKPqAv52g0LAN7AMwAmgzRS+uZSeEXx2f6czN4RHy5uBAKzBjpFp3iHQCE0ZuP4S7nfBLEHFMmAi+8vE2hn1h7+bVwXjwHrvDGUCnjfEEgt+OcZll759CJwB8h94MMGS3GZAgmI5jBQ9tTGeH9EBBIG3Dg4R/YcybAGEAAVK/AQGaAeMClAHzEOgZtg6BPgOOIDBkiQ5eFBXCBFci0phropnQAApZED1z1kSfCfthyKnHdaFsHf0NmGEN6BdAqVVpatsSZmddai92fz94Uijq6pmr6OoYCSirGmvJG3SWS3FE2cBQfT+HlopG4Fsw5agq68iZeSNlpWnBHIedMreuWqGCm1WFrkSSx526WWswAQAA" rules = { JavaScriptDefn(html_escaping.without_repeated_whitespace(""" - window.treescope.handle_copy_click = async (button) => { + this.getRootNode().host.defns.handle_copy_click = async (button) => { const dataToCopy = button.dataset.copy; try { await navigator.clipboard.writeText(dataToCopy); @@ -496,13 +497,25 @@ def html_setup_parts( } }; """)), - CSSStyleRule(html_escaping.without_repeated_whitespace(f""" - @font-face {{ + # Font-face definitions can't live inside a shadow + # DOM node, so we need to inject them into the root document. + JavaScriptDefn(html_escaping.without_repeated_whitespace(""" + if ( + !Array.from(document.fonts.values()).some( + font => font.family == 'Material Symbols Outlined Content Copy"' + ) + ) { + const sheet = new CSSStyleSheet(); + sheet.replaceSync(`@font-face { font-family: 'Material Symbols Outlined Content Copy'; font-style: normal; font-weight: 400; - src: url({font_data_url}) format('woff2'); - }} + src: url({__FONT_DATA_URL__}) format('woff2'); + }`); + document.adoptedStyleSheets = [...document.adoptedStyleSheets, sheet]; + } + """.replace("{__FONT_DATA_URL__}", font_data_url))), + CSSStyleRule(html_escaping.without_repeated_whitespace(f""" {setup_context.collapsed_selector} .copybutton {{ display: none; }} @@ -557,7 +570,8 @@ def render_to_html( repr(self.annotation) ) attributes = html_escaping.without_repeated_whitespace( - "class='copybutton' onClick='treescope.handle_copy_click(this)'" + "class='copybutton'" + " onClick='this.getRootNode().host.defns.handle_copy_click(this)'" f' data-copy="{copy_string_attr}" style="--data-copy:' f" {copy_string_property}; --data-annotation:" f' {annotation_property}" ' @@ -728,6 +742,40 @@ def render_to_text_as_root( return result +TREESCOPE_PREAMBLE_SCRIPT = """(()=> { + const defns = this.getRootNode().host.defns; + let _pendingActions = []; + let _pendingActionHandle = null; + defns.runSoon = (work) => { + const doWork = () => { + const tick = performance.now(); + while (performance.now() - tick < 32) { + if (_pendingActions.length == 0) { + _pendingActionHandle = null; + return; + } else { + const thunk = _pendingActions.shift(); + thunk(); + } + } + _pendingActionHandle = ( + window.requestAnimationFrame(doWork)); + }; + _pendingActions.push(work); + if (_pendingActionHandle === null) { + _pendingActionHandle = ( + window.requestAnimationFrame(doWork)); + } + }; + defns.toggle_root_roundtrip = (rootelt, event) => { + if (event.key == "r") { + rootelt.classList.toggle("roundtrip_mode"); + } + }; +})(); +""" + + def _render_to_html_as_root_streaming( root_node: RenderableTreePart, roundtrip: bool, @@ -776,10 +824,12 @@ def _render_one( stream.write("") if current_js_defns: - stream.write("") + stream.write("") # Render the node itself. node.render_to_html( @@ -800,84 +850,33 @@ def _render_one( background-color: white; color: black; width: fit-content; - margin-left: 2ch; + padding-left: 2ch; line-height: 1.5; + contain: content; + content-visibility: auto; + contain-intrinsic-size: auto none; } """)) stream.write("") # These scripts allow us to defer execution of javascript blocks until after # the content is loaded, avoiding locking up the browser rendering process. - stream.write(html_escaping.without_repeated_whitespace(""" - - """)) + stream.write("") # Render the root node. classnames = "treescope_root" if roundtrip: classnames += " roundtrip_mode" stream.write( - f'
' + f'
' ) _render_one(root_node, True, {}, stream) stream.write("
") - stream.write("") yield stream.getvalue() @@ -907,25 +906,27 @@ def _render_one( ) stream.write("
") - stream.write(" - """)) - stream.write("") + treeroot.replaceWith(treerootClone); + """) + ) + stream.write( + '" + ) yield stream.getvalue() @@ -947,14 +948,9 @@ def render_to_html_as_root( Returns: HTML source for the rendered node. """ - html_src = "".join( - _render_to_html_as_root_streaming(root_node, roundtrip, []) - ) - if compressed: - html_src = html_compression.compress_html( - html_src, include_preamble=True, loading_message="(Loading...)" - ) - return html_src + render_iterator = _render_to_html_as_root_streaming(root_node, roundtrip, []) + html_src = "".join(render_iterator) + return html_encapsulation.encapsulate_html(html_src, compress=compressed) def display_streaming_as_root( @@ -987,26 +983,12 @@ def display_streaming_as_root( render_iterator = _render_to_html_as_root_streaming( root_node, roundtrip, deferreds ) - steal_id = uuid.uuid4().hex - for i, step in enumerate(render_iterator): - if compressed: - if i == 0: - step = html_compression.compress_html( - step, include_preamble=True, loading_message="(Loading...)" - ) - if stealable: - step = f'
{step}
' - else: - step = html_compression.compress_html(step, include_preamble=False) - IPython.display.display(IPython.display.HTML(step)) + encapsulated_iterator = html_encapsulation.encapsulate_streaming_html( + render_iterator, compress=compressed, stealable=stealable + ) - if stealable: - return html_escaping.without_repeated_whitespace( - """
""".replace("{__STEAL_ID__}", steal_id) - ) + for step in encapsulated_iterator: + if step.segment_type == html_encapsulation.SegmentType.FINAL_OUTPUT_STEALER: + return step.html_src + else: + IPython.display.display(IPython.display.HTML(step.html_src)) diff --git a/penzai/treescope/foldable_representation/part_interface.py b/penzai/treescope/foldable_representation/part_interface.py index 06bf75c..d751812 100644 --- a/penzai/treescope/foldable_representation/part_interface.py +++ b/penzai/treescope/foldable_representation/part_interface.py @@ -47,8 +47,8 @@ class JavaScriptDefn: The contents of `source` will be inserted into a - """) - - -def compress_html( - html_src: str, - include_preamble: bool = True, - loading_message: str | None = None, -) -> str: - """Compresses HTML source to an equivalent compressed JavaScript ' - f"{html.escape(loading_message)}" - ) - else: - src_template = ( - '
" - ) - src = src_template.replace( - "__REPLACE_ME_WITH_SERIALIZED__", serialized - ).replace( - "__REPLACE_ME_WITH_UNIQUE_ID__", "compress_html_" + uuid.uuid4().hex - ) - if include_preamble: - src = decompression_preamble() + src - return src diff --git a/penzai/treescope/html_encapsulation.py b/penzai/treescope/html_encapsulation.py new file mode 100644 index 0000000..f1036b5 --- /dev/null +++ b/penzai/treescope/html_encapsulation.py @@ -0,0 +1,369 @@ +# Copyright 2024 The Penzai Authors. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Outer runner to manage interacting HTML outputs in IPython notebooks. + +When rendering HTML outputs in IPython notebooks, we want to support: + +- Compression, so that large repetitive HTML outputs do not cause excessively + large notebook sizes, +- Streaming output, where the structure of an output can be sent to the + browser while NDArrays and other contents are still being loaded and + rendered in the IPython kernel, +- Final output stealing (for streaming output mode), where the temporary + rendering can be moved into the final "result" display cell in IPython + systems that show Out[...] markers, +- Duplication-safety, so that if the same output is repeated multiple times in + a notebook display environment (e.g. in a separate view-only cell), each copy + is rendered independently, +- Shadow-root encapsulation, where the rendering is contained in the "shadow + DOM". This keeps its JS/CSS separate from the main notebook, prevents + interference between different renderings, and may allow rendering + optimizations in some browsers. + +This module defines the necessarly logic to enable these features. It takes as +input a sequence of HTML blobs that should be rendered into a single container +in a streaming fashion, and transforms them into a new sequence of compressed, +duplication-safe HTML blobs that first create the container and then render each +of the contents into it. +""" + +from __future__ import annotations + +import base64 +import dataclasses +import enum +import io +from typing import Iterable, Iterator +import uuid +import zlib + + +class SegmentType(enum.Enum): + """Type of HTML segment.""" + + # The blob that creates and owns the container. + CONTAINER = enum.auto() + # A blob that inserts new data into the initial container, but doesn't display + # any content itself. + CONTAINER_UPDATE = enum.auto() + # A final blob that will take ownership of the container, moving the container + # into itself. Only used when stealing is enabled. + FINAL_OUTPUT_STEALER = enum.auto() + + +@dataclasses.dataclass(frozen=True) +class HTMLOutputSegment: + """A piece of partial HTML output.""" + + html_src: str + segment_type: SegmentType + + +CONTAINER_TEMPLATE = """ + + + +""" + +STEP_TEMPATE = """ +
+ + +
+""" + +COMPRESSED_STEP_TEMPLATE = """ +
+ + +
+""" + +STEALER_TEMPATE = """ + +""" + + +def _prep_html_js_and_strip_comments(src): + stream = io.StringIO() + for line in src.splitlines(): + stripped = line.strip() + if stripped and not stripped.startswith("//"): + stream.write(stripped) + stream.write(" ") + return stream.getvalue()[:-1] + + +def encapsulate_streaming_html( + inner_iterator: Iterable[str], + *, + compress: bool = True, + stealable: bool = False, +) -> Iterator[HTMLOutputSegment]: + """Encapsulates a sequence of inner HTML blobs into robust iframe updates. + + This function accepts an iterator of HTML blobs that should each run in the + same iframe, and transforms them into another iterator of HTML blobs that + can be inserted directly into a notebook environment. The first output will + set up the iframe, and all later updates will insert content into that + original iframe. Optionally, the final output will be a "stealer" that will + move the iframe into itself, ensuring that the iframe is associated with the + correct "result" cell in IPython notebook systems that show Out[...] markers. + Updates will be duplication-safe, in the sense that repeating the same + sequence of outputs in a single HTML page will produce multiple copies of the + iframe, each with the same contents, and the code in the inner iterator will + only be executed once in each iframe. + + Args: + inner_iterator: A nonempty iterator of HTML blobs to encapsulate. + compress: Whether to compress the HTML blobs. + stealable: Whether to include a final "stealer" blob that will move the + iframe into itself. + + Yields: + HTML output segments that can be displayed in a notebook environment or + saved. + """ + inner_iterator = iter(inner_iterator) + + stream = io.StringIO() + # Build the initial iframe, and assign it a unique ID. + # Note: This is unique in the Python program, but if the output is repeated + # multiple times in the notebook output, we may have multiple iframes with + # the same ID. + unique_id_class = f"treescope_out_{uuid.uuid4().hex}" + + outer_content = _prep_html_js_and_strip_comments(CONTAINER_TEMPLATE).replace( + "{__REPLACE_ME_WITH_CONTAINER_ID_CLASS__}", unique_id_class + ) + stream.write(outer_content) + + for i, step_content in enumerate(inner_iterator): + if compress: + # Compress the input string. We use ZLIB, which is natively supported by + # modern browsers. + compressed = zlib.compress( + step_content.encode("utf-8"), zlib.Z_BEST_COMPRESSION + ) + # Serialize it as base64. + serialized = base64.b64encode(compressed).decode("ascii") + # Embed it. + step_content = ( + _prep_html_js_and_strip_comments(COMPRESSED_STEP_TEMPLATE) + .replace("{__REPLACE_ME_WITH_CONTAINER_ID_CLASS__}", unique_id_class) + .replace("{__REPLACE_ME_WITH_STEP__}", str(i)) + .replace("{__REPLACE_ME_WITH_COMPRESSED_CONTENT_HTML__}", serialized) + ) + else: + step_content = ( + _prep_html_js_and_strip_comments(STEP_TEMPATE) + .replace("{__REPLACE_ME_WITH_CONTAINER_ID_CLASS__}", unique_id_class) + .replace("{__REPLACE_ME_WITH_STEP__}", str(i)) + .replace("{__REPLACE_ME_WITH_CONTENT_HTML__}", step_content) + ) + stream.write(step_content) + if i == 0: + segment_type = SegmentType.CONTAINER + else: + segment_type = SegmentType.CONTAINER_UPDATE + yield HTMLOutputSegment( + html_src=stream.getvalue(), segment_type=segment_type + ) + stream = io.StringIO() + + if stealable: + stealer_content = _prep_html_js_and_strip_comments(STEALER_TEMPATE).replace( + "{__REPLACE_ME_WITH_CONTAINER_ID_CLASS__}", unique_id_class + ) + yield HTMLOutputSegment( + html_src=stealer_content, + segment_type=SegmentType.FINAL_OUTPUT_STEALER, + ) + + +def encapsulate_html(html_src: str, compress: bool = True) -> str: + """Encapsulates HTML source code into a duplication-safe container. + + Args: + html_src: The HTML source code to encapsulate. + compress: Whether to compress the HTML source code. + + Returns: + An HTML output segment that can be displayed in a notebook environment or + saved. + """ + [converted] = encapsulate_streaming_html( + [html_src], compress=compress, stealable=False + ) + return converted.html_src diff --git a/penzai/treescope/treescope_ipython.py b/penzai/treescope/treescope_ipython.py index 043e46c..7333c90 100644 --- a/penzai/treescope/treescope_ipython.py +++ b/penzai/treescope/treescope_ipython.py @@ -25,7 +25,6 @@ from penzai.treescope import autovisualize from penzai.treescope import default_renderer from penzai.treescope import figures -from penzai.treescope import html_compression from penzai.treescope import object_inspection from penzai.treescope import selection_rendering from penzai.treescope.arrayviz import array_autovisualizer @@ -68,14 +67,12 @@ def display( raise RuntimeError("Cannot use `display` outside of IPython.") IPython.display.display( IPython.display.HTML( - html_compression.compress_html( - default_renderer.render_to_html( - value, - ignore_exceptions=ignore_exceptions, - roundtrip_mode=roundtrip_mode, - ), - loading_message="(Loading...)", - ) + default_renderer.render_to_html( + value, + ignore_exceptions=ignore_exceptions, + roundtrip_mode=roundtrip_mode, + compressed=True, + ), ) )