Make image as part of the widget model - alternate implementation

ipympl/backend_nbagg.py

@@ -4,10 +4,18 @@
 import json
 from base64 import b64encode
 
+try:
+    from collections.abc import Iterable
+except ImportError:
+    # Python 2.7
+    from collections import Iterable
+
 import matplotlib
+import numpy as np
 from IPython import get_ipython
 from IPython import version_info as ipython_version_info
 from IPython.display import HTML, display
+from ipython_genutils.py3compat import string_types
 from ipywidgets import DOMWidget, widget_serialization
 from matplotlib import is_interactive, rcParams
 from matplotlib._pylab_helpers import Gcf
@@ -18,7 +26,9 @@
     NavigationToolbar2WebAgg,
     TimerTornado,
 )
+from PIL import Image
 from traitlets import (
+    Any,
     Bool,
     CaselessStrEnum,
     CInt,
@@ -142,6 +152,14 @@ class Canvas(DOMWidget, FigureCanvasWebAggCore):
     resizable = Bool(True).tag(sync=True)
     capture_scroll = Bool(False).tag(sync=True)
 
+    # This is a very special widget trait:
+    # We set "sync=True" because we want ipywidgets to consider this
+    # as part of the widget state, but we overwrite send_state so that
+    # it never sync the value with the front-end, the front-end keeps its
+    # own value.
+    # This will still be used by ipywidgets in the case of embedding.
+    _data_url = Any(None).tag(sync=True)
+
     _width = CInt().tag(sync=True)
     _height = CInt().tag(sync=True)
 
@@ -172,12 +190,34 @@ def __init__(self, figure, *args, **kwargs):
 
         self.on_msg(self._handle_message)
 
+        # This will stay True for cases where there is no
+        # front-end (e.g. nbconvert --execute)
+        self.syncing_data_url = True
+
+    # Overwrite ipywidgets's send_state so we don't sync the data_url
+    def send_state(self, key=None):
+        if key is None:
+            keys = self.keys
+        elif isinstance(key, string_types):
+            keys = [key]
+        elif isinstance(key, Iterable):
+            keys = key
+
+        if not self.syncing_data_url:
+            keys = [k for k in keys if k != '_data_url']
+
+        DOMWidget.send_state(self, key=keys)
+
     def _handle_message(self, object, content, buffers):
         # Every content has a "type".
         if content['type'] == 'closing':
             self._closed = True
 
         elif content['type'] == 'initialized':
+            # We stop syncing data url, the front-end is there and
+            # ready to receive diffs
+            self.syncing_data_url = False
+
             _, _, w, h = self.figure.bbox.bounds
             self.manager.resize(w, h)
 
@@ -214,6 +254,19 @@ def send_json(self, content):
             self.send({'data': json.dumps(content)})
 
     def send_binary(self, data):
+        # TODO we should maybe rework the FigureCanvasWebAggCore implementation
+        # so that it has a "refresh" method that we can overwrite
+
+        # Update _data_url
+        if self.syncing_data_url:
+            data = self._last_buff.view(dtype=np.uint8).reshape(
+                (*self._last_buff.shape, 4)
+            )
+            with io.BytesIO() as png:
+                Image.fromarray(data).save(png, format="png")
+                self._data_url = b64encode(png.getvalue()).decode('utf-8')
+
+        # Actually send the data
         self.send({'data': '{"type": "binary"}'}, buffers=[data])
 
     def new_timer(self, *args, **kwargs):
@@ -229,11 +282,11 @@ def _repr_mimebundle_(self, **kwargs):
 
         buf = io.BytesIO()
         self.figure.savefig(buf, format='png', dpi='figure')
-        data_url = b64encode(buf.getvalue()).decode('utf-8')
+        self._data_url = b64encode(buf.getvalue()).decode('utf-8')
 
         data = {
             'text/plain': plaintext,
-            'image/png': data_url,
+            'image/png': self._data_url,
             'application/vnd.jupyter.widget-view+json': {
                 'version_major': 2,
                 'version_minor': 0,

src/mpl_widget.ts

@@ -1,9 +1,11 @@
 import {
     DOMWidgetModel,
     DOMWidgetView,
+    WidgetModel,
     ISerializers,
     unpack_models,
 } from '@jupyter-widgets/base';
+
 import * as utils from './utils';
 
 import { MODULE_VERSION } from './version';
@@ -14,8 +16,9 @@ export class MPLCanvasModel extends DOMWidgetModel {
     requested_size: Array<number> | null;
     resize_requested: boolean;
     ratio: number;
-    waiting: any;
+    waiting_for_image: boolean;
     image: HTMLImageElement;
+
     defaults() {
         return {
             ...super.defaults(),
@@ -32,6 +35,7 @@ export class MPLCanvasModel extends DOMWidgetModel {
             toolbar_position: 'horizontal',
             resizable: true,
             capture_scroll: false,
+            _data_url: null,
             _width: 0,
             _height: 0,
             _figure_label: 'Figure',
@@ -70,6 +74,9 @@ export class MPLCanvasModel extends DOMWidgetModel {
         this.requested_size = null;
         this.resize_requested = false;
         this.ratio = (window.devicePixelRatio || 1) / backingStore;
+
+        this.resize_canvas(this.get('_width'), this.get('_height'));
+
         this._init_image();
 
         this.on('msg:custom', this.on_comm_message.bind(this));
@@ -78,10 +85,30 @@ export class MPLCanvasModel extends DOMWidgetModel {
                 view.update_canvas();
             });
         });
+        this.on('comm_live_update', this.update_disabled.bind(this));
+
+        this.update_disabled();
 
         this.send_initialization_message();
     }
 
+    get disabled(): boolean {
+        return !this.comm_live;
+    }
+
+    update_disabled(): void {
+        this.set('resizable', !this.disabled);
+    }
+
+    sync(method: string, model: WidgetModel, options: any = {}) {
+        // Make sure we don't sync the data_url, we don't need it to be synced
+        if (options.attrs) {
+            delete options.attrs['_data_url'];
+        }
+
+        super.sync(method, model, options);
+    }
+
     send_message(type: string, message: { [index: string]: any } = {}) {
         message['type'] = type;
 
@@ -96,15 +123,15 @@ export class MPLCanvasModel extends DOMWidgetModel {
             });
         }
 
-        this.send_message('send_image_mode');
         this.send_message('refresh');
+        this.send_message('send_image_mode');
 
         this.send_message('initialized');
     }
 
     send_draw_message() {
-        if (!this.waiting) {
-            this.waiting = true;
+        if (!this.waiting_for_image) {
+            this.waiting_for_image = true;
             this.send_message('draw');
         }
     }
@@ -160,8 +187,8 @@ export class MPLCanvasModel extends DOMWidgetModel {
     }
 
     resize_canvas(width: number, height: number) {
-        this.offscreen_canvas.setAttribute('width', `${width * this.ratio}`);
-        this.offscreen_canvas.setAttribute('height', `${height * this.ratio}`);
+        this.offscreen_canvas.width = width * this.ratio;
+        this.offscreen_canvas.height = height * this.ratio;
     }
 
     handle_rubberband(msg: any) {
@@ -204,7 +231,9 @@ export class MPLCanvasModel extends DOMWidgetModel {
 
         this.image.src = image_url;
 
-        this.waiting = false;
+        this.set('_data_url', this.offscreen_canvas.toDataURL());
+
+        this.waiting_for_image = false;
     }
 
     handle_history_buttons(msg: any) {
@@ -240,7 +269,8 @@ export class MPLCanvasModel extends DOMWidgetModel {
     }
 
     _init_image() {
-        this.image = document.createElement('img');
+        this.image = new Image();
+
         this.image.onload = () => {
             if (this.get('_image_mode') === 'full') {
                 // Full images could contain transparency (where diff images
@@ -259,6 +289,12 @@ export class MPLCanvasModel extends DOMWidgetModel {
                 view.update_canvas();
             });
         };
+
+        const dataUrl = this.get('_data_url');
+
+        if (dataUrl !== null) {
+            this.image.src = dataUrl;
+        }
     }
 
     _for_each_view(callback: any) {
@@ -285,12 +321,12 @@ export class MPLCanvasView extends DOMWidgetView {
     context: CanvasRenderingContext2D;
     top_canvas: HTMLCanvasElement;
     top_context: CanvasRenderingContext2D;
-    waiting: boolean;
     footer: HTMLDivElement;
     model: MPLCanvasModel;
     private _key: string | null;
     private _resize_event: (event: MouseEvent) => void;
     private _stop_resize_event: () => void;
+
     render() {
         this.resizing = false;
         this.resize_handle_size = 20;
@@ -313,8 +349,6 @@ export class MPLCanvasView extends DOMWidgetView {
         window.addEventListener('mousemove', this._resize_event);
         window.addEventListener('mouseup', this._stop_resize_event);
 
-        this.waiting = false;
-
         return this.create_child_view(this.model.get('toolbar')).then(
             (toolbar_view) => {
                 this.toolbar_view = toolbar_view;

src/toolbar_widget.ts

@@ -25,6 +25,13 @@ export class ToolbarView extends DOMWidgetView {
     toggle_button: HTMLButtonElement;
     toolbar: HTMLDivElement;
     buttons: { [index: string]: HTMLButtonElement };
+
+    initialize(parameters: any) {
+        super.initialize(parameters);
+
+        this.on('comm_live_update', this.update_disabled.bind(this));
+    }
+
     render(): void {
         this.el.classList.add(
             'jupyter-widgets',
@@ -100,6 +107,20 @@ export class ToolbarView extends DOMWidgetView {
         this.set_orientation(this.el);
         this.set_orientation(this.toolbar);
         this.set_buttons_style();
+
+        this.update_disabled();
+    }
+
+    get disabled(): boolean {
+        return !this.model.comm_live;
+    }
+
+    update_disabled(): void {
+        // Disable buttons
+        this.toggle_button.disabled = this.disabled;
+        if (this.disabled) {
+            this.toolbar.style.display = 'none';
+        }
     }
 
     set_orientation(el: HTMLElement): void {