Snapshot functionality (#37)

* include snapshot in repr mimetypes so you can export notebooks and see an image

* better implementation

* add tests and move some code around

* Add docs

* remove empty line

Co-authored-by: David Hoese <[email protected]>

* better link

* Make Snapshot class a subclass of DisplayObject.

* tweak comment

* fix tests

Co-authored-by: David Hoese <[email protected]>

docs/examples/index.rst

@@ -9,8 +9,12 @@ This is a list of example notebooks demonstrating the use of jupyter_rfb:
     hello_world.ipynb
     interaction1.ipynb
     interaction2.ipynb
+    ipywidgets_embed.ipynb
     performance.ipynb
     push.ipynb
+    pygfx1.ipynb
+    pygfx2.ipynb
+    pygfx3.ipynb
     video.ipynb
     vispy1.ipynb
     wgpu1.ipynb

docs/guide.rst

@@ -86,6 +86,22 @@ second. This is to avoid spamming the io and server process. The
 throttling applies to the resize, scroll, and pointer_move events.
 
 
+Taking snapshots
+----------------
+
+In a notebook, the :meth:`.snapshot() <jupyter_rfb.RemoteFrameBuffer.snapshot>`
+method can be used to create a picture of the current state of the
+widget. This image remains visible when the notebook is in off-line
+mode (e.g. in nbviewer). This functionality can be convenient if you're
+using a notebook to tell a story, and you want to display a certain
+result that is also visible in off-line mode.
+
+When a widget is first displayed, it also automatically creates a
+snapshot, which is hidden by default, but is visible when the
+widget itself is not loaded. In other words, example notebooks
+have pretty pictures!
+
+
 Exceptions and logging
 ----------------------
 
@@ -109,7 +125,7 @@ Measuring statistics
 The ``RemoteFrameBuffer`` class has a method ``get_stats()`` that
 returns a dict with performance metrics:
 
-.. code-block::
+.. code-block:: py
 
     >>> w.reset_stats()  # start measuring
         ... interact or run a test

js/lib/widget.js

@@ -133,6 +133,11 @@ var RemoteFrameBufferView = widgets.DOMWidgetView.extend({
     // Defines how the widget gets rendered into the DOM
     render: function () {
         var that = this;
+
+        // Hide initial snapshot
+        for (let el of document.getElementsByClassName("initial-snapshot-" + this.model.model_id)) {
+            el.style.display = "none";
+        }
 
         // Create a stub element that can grab focus
         this.focus_el = document.createElement("a");

jupyter_rfb/_utils.py

@@ -0,0 +1,93 @@
+import io
+import builtins
+import traceback
+from base64 import encodebytes
+
+from IPython.display import DisplayObject
+import ipywidgets
+
+from ._png import array2png
+
+_original_print = builtins.print
+
+
+class RFBOutputContext(ipywidgets.Output):
+    """An output widget with a different implementation of the context manager.
+
+    Handles prints and errors in a more reliable way, that is also
+    lightweight (i.e. no peformance cost).
+
+    See https://github.com/vispy/jupyter_rfb/issues/35
+    """
+
+    capture_print = False
+    _prev_print = None
+
+    def print(self, *args, **kwargs):
+        """Print function that show up in the output."""
+        f = io.StringIO()
+        kwargs.pop("file", None)
+        _original_print(*args, file=f, flush=True, **kwargs)
+        text = f.getvalue()
+        self.append_stdout(text)
+
+    def __enter__(self):
+        """Enter context, replace print function."""
+        if self.capture_print:
+            self._prev_print = builtins.print
+            builtins.print = self.print
+        return self
+
+    def __exit__(self, etype, value, tb):
+        """Exit context, restore print function and show any errors."""
+        if self.capture_print and self._prev_print is not None:
+            builtins.print = self._prev_print
+            self._prev_print = None
+        if etype:
+            err = "".join(traceback.format_exception(etype, value, tb))
+            self.append_stderr(err)
+            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
+        png_data = array2png(self.data)
+        preamble = "data:image/png;base64,"
+        src = preamble + encodebytes(png_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:1px 3px; "
+        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()

jupyter_rfb/widget.py

@@ -12,9 +12,6 @@
 the server will slow down too.
 """
 
-import io
-import builtins
-import traceback
 import asyncio
 import time
 from base64 import encodebytes
@@ -25,47 +22,7 @@
 from traitlets import Bool, Dict, Int, Unicode
 
 from ._png import array2png
-
-
-_original_print = builtins.print
-
-
-class OutputContext(ipywidgets.Output):
-    """An output widget with a different implementation of the context manager.
-
-    Handles prints and errors in a more reliable way, that is also
-    lightweight (i.e. no peformance cost).
-
-    See https://github.com/vispy/jupyter_rfb/issues/35
-    """
-
-    capture_print = False
-    _prev_print = None
-
-    def print(self, *args, **kwargs):
-        """Print function that show up in the output."""
-        f = io.StringIO()
-        kwargs.pop("file", None)
-        _original_print(*args, file=f, flush=True, **kwargs)
-        text = f.getvalue()
-        self.append_stdout(text)
-
-    def __enter__(self):
-        """Enter context, replace print function."""
-        if self.capture_print:
-            self._prev_print = builtins.print
-            builtins.print = self.print
-        return self
-
-    def __exit__(self, etype, value, tb):
-        """Exit context, restore print function and show any errors."""
-        if self.capture_print and self._prev_print is not None:
-            builtins.print = self._prev_print
-            self._prev_print = None
-        if etype:
-            err = "".join(traceback.format_exception(etype, value, tb))
-            self.append_stderr(err)
-            return True  # declare that we handled the exception
+from ._utils import RFBOutputContext, Snapshot
 
 
 @ipywidgets.register
@@ -113,21 +70,59 @@ class RemoteFrameBuffer(ipywidgets.DOMWidget):
 
     def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
-        # Setup an output widget, so that any prints and errors in our
-        # callbacks are actually shown. We display the output in the cell-output
+        self._ipython_display_ = None  # we use _repr_mimebundle_ instread
+        # Setup an output widget, so that any errors in our callbacks
+        # are actually shown. We display the output in the cell-output
         # corresponding to the cell that instantiates the widget.
-        self._output_context = OutputContext()
+        self._output_context = RFBOutputContext()
         display(self._output_context)
         # Init attributes for drawing
         self._rfb_draw_requested = False
         self._rfb_frame_index = 0
         self._rfb_last_confirmed_index = 0
+        self._rfb_last_resize_event = None
         # Init stats
         self.reset_stats()
         # Setup events
         self.on_msg(self._rfb_handle_msg)
         self.observe(self._rfb_schedule_maybe_draw, names=["frame_feedback"])
 
+    def _repr_mimebundle_(self, **kwargs):
+        data = {}
+
+        # Always add plain text
+        plaintext = repr(self)
+        if len(plaintext) > 110:
+            plaintext = plaintext[:110] + "…"
+        data["text/plain"] = plaintext
+
+        # Get the actual representation
+        try:
+            data.update(super()._repr_mimebundle_(**kwargs))
+        except Exception:
+            # On 7.6.3 and below, _ipython_display_ is used instead of _repr_mimebundle_.
+            # We fill in the widget representation that has been in use for 5+ years.
+            data["application/vnd.jupyter.widget-view+json"] = {
+                "version_major": 2,
+                "version_minor": 0,
+                "model_id": self._model_id,
+            }
+
+        # Add initial snapshot. It would be awesome if, when the
+        # notebook is offline, this representation is used instead of
+        # application/vnd.jupyter.widget-view+json. And in fact, Gihub's
+        # renderer does this. Unfortunately, nbconvert still selects
+        # the widget mimetype.
+        # So instead, we display() the snapshot right in front of the
+        # actual widget view, and when the widget view is created, it
+        # hides the snapshot. Ha! That way, the snapshot is
+        # automatically shown when the widget is not loaded!
+        if self._view_name is not None:
+            # data["text/html"] = self.snapshot()._get_html()
+            display(self.snapshot(None, _initial=True))
+
+        return data
+
     def print(self, *args, **kwargs):
         """Print to the widget's output area (For debugging purposes).
 
@@ -150,12 +145,58 @@ def _rfb_handle_msg(self, widget, content, buffers):
         """Receive custom messages and filter our events."""
         if "event_type" in content:
             if content["event_type"] == "resize":
+                self._rfb_last_resize_event = content
                 self.request_draw()
+            elif content["event_type"] == "close":
+                self._repr_mimebundle_ = None
             with self._output_context:
                 self.handle_event(content)
 
     # ---- drawing
 
+    def snapshot(self, pixel_ratio=None, _initial=False):
+        """Create a snapshot of the current state of the widget.
+
+        Returns an ``IPython DisplayObject`` that can simply be used as
+        a cell output. May also return None if ``get_frame()`` produces
+        None. The display object has a ``data`` attribute that holds
+        the image array data (typically a numpy array).
+        """
+        # Start with a resize event to the appropriate pixel ratio
+        ref_resize_event = self._rfb_last_resize_event
+        if ref_resize_event:
+            w = ref_resize_event["width"]
+            h = ref_resize_event["height"]
+        else:
+            pixel_ratio = pixel_ratio or 1
+            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
+        if pixel_ratio:
+            evt = {
+                "event_type": "resize",
+                "width": w,
+                "height": h,
+                "pixel_ratio": pixel_ratio,
+            }
+            self.handle_event(evt)
+        # Render a frame
+        array = self.get_frame()
+        # Reset pixel ratio
+        if ref_resize_event and pixel_ratio:
+            self.handle_event(ref_resize_event)
+        # Create snapshot object
+        if array is None:
+            return None
+        elif _initial:
+            title = "initial snapshot"
+            class_name = "initial-snapshot-" + self._model_id
+        else:
+            title = "snapshot"
+            class_name = "snapshot-" + self._model_id
+
+        return Snapshot(array, w, h, title, class_name)
+
     def request_draw(self):
         """Schedule a new draw when the widget is ready for it.
 

setup.cfg

@@ -4,7 +4,7 @@ universal=1
 [flake8]
 max-line-length = 120
 ignore = D107, D202, W503
-exclude = .git,__pycache__,build,dist
+exclude = .git,__pycache__,build,dist,.ipynb_checkpoints
 
 [isort]
 sections = FUTURE,STDLIB,THIRDPARTY,FIRSTPARTY,LOCALFOLDER