Refactor how snapshot works (#152)

* WIP refactor how snapshot works

* clean up

* fix tests

Author: almarklein

jupyter_rfb/_utils.py

@@ -1,9 +1,7 @@
 import io
 import builtins
 import traceback
-from base64 import encodebytes
 
-from IPython.display import DisplayObject
 import ipywidgets
 
 from ._png import array2png
@@ -75,73 +73,6 @@ def __exit__(self, etype, value, tb):
             return True  # declare that we handled the exception
 
 
-class Snapshot(DisplayObject):
-    """An IPython DisplayObject representing an image snapshot.
-
-    The ``data`` attribute is the image array object. One could use
-    this to process the data further, e.g. storing it to disk.
-    """
-
-    # Not an IPython.display.Image, because we want to use some HTML to
-    # give it a custom css class and a title.
-
-    def __init__(self, data, width, height, title="snapshot", class_name=None):
-        super().__init__(data)
-        self.width = width
-        self.height = height
-        self.title = title
-        self.class_name = class_name
-
-    def _check_data(self):
-        assert hasattr(self.data, "shape") and hasattr(self.data, "dtype")
-
-    def _repr_mimebundle_(self, **kwargs):
-        return {"text/html": self._repr_html_()}
-
-    def _repr_html_(self):
-        # Convert to PNG
-        mimetype, data = array2compressed(self.data, 70)
-        src = f"data:image/{mimetype};base64," + encodebytes(data).decode()
-        # Create html repr
-        class_str = f"class='{self.class_name}'" if self.class_name else ""
-        img_style = f"width:{self.width}px;height:{self.height}px;"
-        tt_style = "position: absolute; top:0; left:0; padding:3px 4px; border-radius:0 0 4px 0;"
-        tt_style += (
-            "background: #777; color:#fff; font-size: 90%; font-family:sans-serif; "
-        )
-        html = f"""
-            <div {class_str} style='position:relative;'>
-                <img src='{src}' style='{img_style}' />
-                <div style='{tt_style}'>{self.title}</div>
-            </div>
-            """
-        return html.replace("\n", "").replace("    ", "").strip()
-
-
 def remove_rfb_models_from_nb(d):
-    """Remove the widget model output from a notebook dict.
-
-    Given a notebook as a dict (loaded using json), remove the widget
-    model output if there is also a text/html snapshot output.
-
-    This is to work around the fact that nbsphinx favors the model over
-    the text/html output. Which is sad, because that's where we put the
-    initial screenshot for offline viewing.
-    """
-
-    to_remove = set()
-    for key, val in d.items():
-        if key == "cells" and isinstance(val, list):
-            for v in val:
-                remove_rfb_models_from_nb(v)
-        elif key == "outputs" and isinstance(val, list):
-            for v in val:
-                data = v.get("data", None)
-                if data:
-                    remove_rfb_models_from_nb(data)
-        elif key == "application/vnd.jupyter.widget-view+json":
-            html_sibling = d.get("text/html", [])
-            if html_sibling and "<div class='snapshot-" in html_sibling[0]:
-                to_remove.add(key)
-    for key in to_remove:
-        d.pop(key)
+    """Deprecated, is a no-op for backwards compatibility."""
+    pass

jupyter_rfb/widget.py

@@ -16,11 +16,10 @@
 from importlib.resources import files as resource_files
 
 import anywidget
-import numpy as np
-from IPython.display import display
+from IPython.display import display, HTML
 from traitlets import Bool, Dict, Int, Unicode
 
-from ._utils import array2compressed, RFBOutputContext, Snapshot
+from ._utils import array2compressed, RFBOutputContext
 
 
 def _load_js_and_css():
@@ -105,7 +104,7 @@ def __init__(self, *args, **kwargs):
         display(self._output_context)
         # Init attributes for drawing
         self._rfb_last_frame = None
-        self._rfb_pending_display = None
+        self._rfb_pending_snapshot_display = None
         self._rfb_draw_requested = False
         self._rfb_frame_index = 0
         self._rfb_last_confirmed_index = 0
@@ -122,21 +121,6 @@ def __init__(self, *args, **kwargs):
             names=["_frame_feedback", "_has_visible_views"],
         )
 
-    def _repr_mimebundle_(self, **kwargs):
-        # Use default
-        result = anywidget.AnyWidget._repr_mimebundle_(self, **kwargs)
-        # Get dict to add more data
-        data = None
-        if isinstance(result, tuple):
-            data = result[0]
-        elif isinstance(result, dict):
-            data = result
-        # Add initial snapshot if we have it
-        if data and self._rfb_pending_display and self._rfb_last_frame is not None:
-            data["text/html"] = self.snapshot()._repr_html_()
-            self._rfb_pending_display = None  # no need to reload
-        return result
-
     def print(self, *args, **kwargs):
         """Print to the widget's output area (for debugging purposes).
 
@@ -155,7 +139,7 @@ def close(self, *args, **kwargs):
         # When the widget is closed, we notify by creating a close event. The
         # same event is emitted from JS when the model is closed in the client.
         anywidget.AnyWidget.close(self, *args, **kwargs)
-        self._rfb_handle_msg(self, {"type": "close", "event_type": "close"}, [])
+        self._rfb_handle_msg(self, {"type": "close"}, [])
 
     def _rfb_handle_msg(self, widget, content, buffers):
         """Receive custom messages and filter our events."""
@@ -193,43 +177,60 @@ def _rfb_handle_msg(self, widget, content, buffers):
 
     # ---- drawing
 
-    def display(self):
-        """Display the widget.
-
-        The benefit of using this (instead of using the widget as the last expression of a cell)
-        is that an html snapshot is added to the output. This happens either directly,
-        or on the next drawn frame (if no frames have been drawn yet).
-        """
-        self._rfb_pending_display = True
-        id = display(self, display_id=True)
-        if self._rfb_pending_display:  # i.e. is not set to None by __repr__
-            self._rfb_pending_display = id
-
     def snapshot(self, pixel_ratio=None):
-        """Create a snapshot of the current state of the widget.
+        """Render a frame and include the resulting image in the output.
 
-        Returns an ``IPython DisplayObject`` that can simply be used as
-        a cell output. The display object has a ``data`` attribute that holds
-        the image array data (typically a numpy array).
+        An initial placeholder output is produced, which is replaced by an html
+        ``<img>`` as soon as the next frame is rendered.
 
-        The ``pixel_ratio`` argument is deprecated and ignored.
+        If the widget is not displayed yet, a resize event is emitted to mimic a widget
+        size. This happens at most once in the widget's lifetime. It will use the
+        ``css_width`` and ``css_height`` if they are expressed in ``px``, and otherwise
+        default to 500 or 300 pixels respectively. The ``pixel_ratio`` argument is then
+        used to calculate the physical size.
         """
-        # Get the current size
-        ref_resize_event = self._rfb_last_resize_event
-        if ref_resize_event:
-            w = ref_resize_event["width"]
-            h = ref_resize_event["height"]
-        else:
+        if self._rfb_last_resize_event is None:
             css_width, css_height = self.css_width, self.css_height
             w = float(css_width[:-2]) if css_width.endswith("px") else 500
             h = float(css_height[:-2]) if css_height.endswith("px") else 300
-        # Get last frame or single-pixel image
-        array = self._rfb_last_frame
-        if array is None:
-            array = np.ones((1, 1, 3), np.uint8) * 127
-        # Super-weird, but it looks like nbsphinx only selects the text/html field when we use a css class
-        # that starts with 'snapshot-'. Is this some upstream hack to make jupyter-rfb work, that we don't know of?
-        return Snapshot(array, w, h, "snapshot", f"snapshot-rfb model{self._model_id}")
+            r = float(pixel_ratio) if pixel_ratio is not None else 1.0
+            pw, ph = int(w * r), (h * r)
+            event = {
+                "type": "resize",
+                "width": pw / r,
+                "height": ph / r,
+                "pwidth": pw,
+                "pheight": ph,
+                "ratio": r,
+                "timestamp": 0,
+            }
+            self._rfb_handle_msg(self, event, [])
+
+        self._rfb_pending_snapshot_display = display(
+            HTML(
+                "<div style='display: inline-block; padding: 5px; border-radius: 5px; background:#ddd; color:#000'>pending screenshot ...</span>"
+            ),
+            display_id=True,
+        )
+        self.request_draw()
+
+        # Note: It could be that _replace_snapshot() is called directly,
+        # (and _rfb_pending_snapshot_display set to None). But it could
+        # also be that it is called later. In any case, we just return None
+
+    def _replace_snapshot(self, array):
+        pending_display = self._rfb_pending_snapshot_display
+        self._rfb_pending_snapshot_display = None
+
+        event = self._rfb_last_resize_event or {}
+        w = event.get("width", array.shape[1])
+        h = event.get("height", array.shape[0])
+
+        mimetype, data = array2compressed(array, 70)
+        src = f"data:image/{mimetype};base64," + encodebytes(data).decode()
+        html = f"<img src='{src}' style='width:{w}px;height:{h}px;' />"
+
+        pending_display.update(HTML(html))
 
     def request_draw(self):
         """Schedule a new draw. This method itself returns immediately.
@@ -279,7 +280,7 @@ def _rfb_maybe_draw(self):
         should_draw = (
             self._rfb_draw_requested
             and frames_in_flight < self.max_buffered_frames
-            and self._has_visible_views
+            and (self._has_visible_views or self._rfb_pending_snapshot_display)
         )
         # Do the draw if we should.
         if should_draw:
@@ -359,9 +360,8 @@ def _rfb_send_frame(self, array, is_lossless_redraw=False):
                 self._rfb_last_confirmed_index = self._rfb_frame_index - 1
 
         # Reload the output if we did not have a frame when the widget was first loaded
-        if self._rfb_pending_display is not None:
-            if self._rfb_last_resize_event is not None:
-                self._rfb_pending_display.update(self)  # -> calls _repr_mimebundle_
+        if self._rfb_pending_snapshot_display is not None:
+            self._replace_snapshot(array)
 
         # Compose message and send
         msg = dict(

tests/test_utils.py

@@ -2,7 +2,7 @@
 
 import pytest
 import numpy as np
-from jupyter_rfb._utils import array2compressed, RFBOutputContext, Snapshot
+from jupyter_rfb._utils import array2compressed, RFBOutputContext
 from jupyter_rfb import _jpg
 
 
@@ -113,23 +113,3 @@ def test_output_context():
     # The print is a proper print
     c.print("foo", "bar", sep="-", end=".")
     assert c.stdouts[-1] == "foo-bar."
-
-
-def test_snapshot():
-    """Test the Snapshot class."""
-
-    a = np.zeros((10, 10), np.uint8)
-
-    s = Snapshot(a, 5, 5, "footitle", "KLS")
-
-    # The get_array method returns the raw data
-    assert s.data is a
-
-    # Most importantly, it has a Jupyter mime data!
-    data = s._repr_mimebundle_()
-    assert "text/html" in data
-    html = data["text/html"]
-    assert "data:image/" in html  # looks like the png/jpg is in there
-    assert "width:5px" in html and "height:5px" in html  # logical size
-    assert "class='KLS'" in html  # css class name
-    assert "footitle" in html  # the title

tests/test_widget.py

@@ -9,7 +9,6 @@
 import numpy as np
 from pytest import raises
 from jupyter_rfb import RemoteFrameBuffer
-from jupyter_rfb._utils import Snapshot
 from traitlets import TraitError
 
 
@@ -300,18 +299,10 @@ def test_print():
 
 
 def test_snapshot():
-    """Test that the widget has a snapshot method that produces a Snapshot."""
+    """Test that the widget has a snapshot method."""
     w = MyRFB()
     s = w.snapshot()
-    assert isinstance(s, Snapshot)
-    assert s.data.shape == (1, 1, 3)
-
-    w.request_draw()
-    w._rfb_maybe_draw()  # similate first frame
-
-    s = w.snapshot()
-    assert isinstance(s, Snapshot)
-    assert np.all(s.data == w.get_frame())
+    assert s is None  # snapshot() uses display()
 
 
 def test_use_websocket():