Merge pull request #1194 from maartenbreddels/nested-buffers

Any descendant of the state can be a binary_type

Author: jasongrout

ipywidgets/widgets/tests/test_traits.py

@@ -2,11 +2,13 @@
 # Distributed under the terms of the Modified BSD License.
 
 """Test trait types of the widget packages."""
+import array
 
 from unittest import TestCase
 from traitlets import HasTraits
 from traitlets.tests.test_traitlets import TraitTestBase
 from ipywidgets import Color
+from ipywidgets.widgets.widget import _remove_buffers, _put_buffers
 
 
 class ColorTrait(HasTraits):
@@ -19,3 +21,53 @@ class TestColor(TraitTestBase):
     _good_values = ["blue", "#AA0", "#FFFFFF"]
     _bad_values = ["vanilla", "blues"]
 
+
+class TestBuffers(TestCase):
+    def test_remove_and_put_buffers(self):
+        mv1 =  memoryview(b'test1')
+        mv2 =  memoryview(b'test2')
+        state = {'plain': [0, 'text'], # should not get removed
+                 'x': {'ar': mv1}, # should result in an empty dict
+                 'y': {'shape': (10, 10), 'data': mv1},
+                 'z': (mv1, mv2), # tests tuple assigment
+                 'top': mv1, # test a top level removal
+                 'deep': {'a': 1, 'b':[0,{'deeper':mv2}]}} # deeply nested
+        plain = state['plain']
+        x = state['x']
+        y = state['y']
+        y_shape = y['shape']
+        state_before = state
+        state, buffer_paths, buffers = _remove_buffers(state)
+
+        # check if buffers are removed
+        self.assertIn('plain', state)
+        self.assertIn('shape', state['y'])
+        self.assertNotIn('ar', state['x'])
+        self.assertEqual(state['x'], {})
+        self.assertNotIn('data', state['y'])
+        self.assertNotIn(mv1, state['z'])
+        self.assertNotIn(mv1, state['z'])
+        self.assertNotIn('top', state)
+        self.assertIn('deep', state)
+        self.assertIn('b', state['deep'])
+        self.assertNotIn('deeper', state['deep']['b'][1])
+
+        # check that items that didn't need change aren't touched
+        self.assertIsNot(state, state_before)
+        self.assertIs(state['plain'], plain)
+        self.assertIsNot(state['x'], x)
+        self.assertIsNot(state['y'], y)
+        self.assertIs(state['y']['shape'], y_shape)
+
+        # check that the buffer paths really point to the right buffer
+        for path, buffer in [(['x', 'ar'], mv1), (['y', 'data'], mv1), (['z', 0], mv1), (['z', 1], mv2),\
+                             (['top'], mv1), (['deep', 'b', 1, 'deeper'], mv2)]:
+            self.assertIn(path, buffer_paths, "%r not in path" % path)
+            index = buffer_paths.index(path)
+            self.assertEqual(buffer, buffers[index])
+
+        # and check that we can put it back together again
+        _put_buffers(state, buffer_paths, buffers)
+        # we know that tuples get converted to list, so help the comparison by changing the tuple to a list
+        state_before['z'] = list(state_before['z'])
+        self.assertEqual(state_before, state)

ipywidgets/widgets/widget.py

@@ -50,6 +50,78 @@ def _json_to_widget(x, obj):
 else:
     _binary_types = (memoryview, buffer)
 
+def _put_buffers(state, buffer_paths, buffers):
+    """The inverse of _remove_buffers, except here we modify the existing dict/lists.
+    Modifying should be fine, since this is used when state comes from the wire.
+    """
+    for buffer_path, buffer in zip(buffer_paths, buffers):
+        # we'd like to set say sync_data['x'][0]['y'] = buffer
+        # where buffer_path in this example would be ['x', 0, 'y']
+        obj = state
+        for key in buffer_path[:-1]:
+            obj = obj[key]
+        obj[buffer_path[-1]] = buffer
+
+def _separate_buffers(substate, path, buffer_paths, buffers):
+    """For internal, see _remove_buffers"""
+    # remove binary types from dicts and lists, but keep track of their paths
+    # any part of the dict/list that needs modification will be cloned, so the original stays untouched
+    # e.g. {'x': {'ar': ar}, 'y': [ar2, ar3]}, where ar/ar2/ar3 are binary types
+    # will result in {'x': {}, 'y': [None, None]}, [ar, ar2, ar3], [['x', 'ar'], ['y', 0], ['y', 1]]
+    # instead of removing elements from the list, this will make replacing the buffers on the js side much easier
+    if isinstance(substate, (list, tuple)):
+        is_cloned = False
+        for i, v in enumerate(substate):
+            if isinstance(v, _binary_types):
+                if not is_cloned:
+                    substate = list(substate) # shallow clone list/tuple
+                    is_cloned = True
+                substate[i] = None
+                buffers.append(v)
+                buffer_paths.append(path + [i])
+            elif isinstance(v, (dict, list, tuple)):
+                vnew = _separate_buffers(v, path + [i], buffer_paths, buffers)
+                if v is not vnew: # only assign when value changed
+                    if not is_cloned:
+                        substate = list(substate) # clone list/tuple
+                        is_cloned = True
+                    substate[i] = vnew
+    elif isinstance(substate, dict):
+        is_cloned = False
+        for k, v in substate.items():
+            if isinstance(v, _binary_types):
+                if not is_cloned:
+                    substate = dict(substate) # shallow clone dict
+                    is_cloned = True
+                del substate[k]
+                buffers.append(v)
+                buffer_paths.append(path + [k])
+            elif isinstance(v, (dict, list, tuple)):
+                vnew = _separate_buffers(v, path + [k], buffer_paths, buffers)
+                if v is not vnew: # only assign when value changed
+                    if not is_cloned:
+                        substate = dict(substate) # clone list/tuple
+                        is_cloned = True
+                    substate[k] = vnew
+    else:
+        raise ValueError("expected state to be a list or dict, not %r" % substate)
+    return substate
+
+
+def _remove_buffers(state):
+    """Return (state_without_buffers, buffer_paths, buffers) for binary message parts
+
+    As an example:
+    >>> state = {'plain': [0, 'text'], 'x': {'ar': memoryview(ar1)}, 'y': {'shape': (10,10), 'data': memoryview(ar2)}}
+    >>> _remove_buffers(state)
+    ({'plain': [0, 'text']}, {'x': {}, 'y': {'shape': (10, 10)}}, [['x', 'ar'], ['y', 'data']],
+     [<memory at 0x107ffec48>, <memory at 0x107ffed08>])
+    """
+    buffer_paths, buffers = [], []
+    state = _separate_buffers(state, [], buffer_paths, buffers)
+    return state, buffer_paths, buffers
+
+
 class CallbackDispatcher(LoggingConfigurable):
     """A structure for registering and running callbacks"""
     callbacks = List()
@@ -213,17 +285,15 @@ def __del__(self):
     def open(self):
         """Open a comm to the frontend if one isn't already open."""
         if self.comm is None:
-            state, buffer_keys, buffers = self._split_state_buffers(self.get_state())
+            state, buffer_paths, buffers = _remove_buffers(self.get_state())
 
-            args = dict(target_name='jupyter.widget', data=state)
+            args = dict(target_name='jupyter.widget',
+                        data={'state': state, 'buffer_paths': buffer_paths},
+                        buffers=buffers)
             if self._model_id is not None:
                 args['comm_id'] = self._model_id
 
             self.comm = Comm(**args)
-            if buffers:
-                # FIXME: workaround ipykernel missing binary message support in open-on-init
-                # send state with binary elements as second message
-                self.send_state()
 
     @observe('comm')
     def _comm_changed(self, change):
@@ -258,16 +328,6 @@ def close(self):
             self.comm = None
             self._ipython_display_ = None
 
-    def _split_state_buffers(self, state):
-        """Return (state_without_buffers, buffer_keys, buffers) for binary message parts"""
-        buffer_keys, buffers = [], []
-        for k, v in list(state.items()):
-            if isinstance(v, _binary_types):
-                state.pop(k)
-                buffers.append(v)
-                buffer_keys.append(k)
-        return state, buffer_keys, buffers
-
     def send_state(self, key=None):
         """Sends the widget state, or a piece of it, to the front-end.
 
@@ -277,8 +337,8 @@ def send_state(self, key=None):
             A single property's name or iterable of property names to sync with the front-end.
         """
         state = self.get_state(key=key)
-        state, buffer_keys, buffers = self._split_state_buffers(state)
-        msg = {'method': 'update', 'state': state, 'buffers': buffer_keys}
+        state, buffer_paths, buffers = _remove_buffers(state)
+        msg = {'method': 'update', 'state': state, 'buffer_paths': buffer_paths}
         self._send(msg, buffers=buffers)
 
     def get_state(self, key=None, drop_defaults=False):
@@ -453,8 +513,8 @@ def _handle_msg(self, msg):
             if 'sync_data' in data:
                 # get binary buffers too
                 sync_data = data['sync_data']
-                for i,k in enumerate(data.get('buffer_keys', [])):
-                    sync_data[k] = msg['buffers'][i]
+                if 'buffer_paths' in data:
+                    _put_buffers(sync_data, data['buffer_paths'], msg['buffers'])
                 self.set_state(sync_data) # handles all methods
 
         # Handle a state request.

jupyter-js-widgets/src/manager-base.ts

@@ -185,12 +185,14 @@ abstract class ManagerBase<T> {
      * Handle when a comm is opened.
      */
     handle_comm_open(comm: shims.services.Comm, msg: services.KernelMessage.ICommOpenMsg): Promise<Backbone.Model> {
+        var data = (msg.content.data as any);
+        utils.put_buffers(data.state, data.buffer_paths, msg.buffers)
         return this.new_model({
-            model_name: msg.content.data['_model_name'],
-            model_module: msg.content.data['_model_module'],
-            model_module_version: msg.content.data['_model_module_version'],
+            model_name: data.state['_model_name'],
+            model_module: data.state['_model_module'],
+            model_module_version: data.state['_model_module_version'],
             comm: comm
-        }, msg.content.data).catch(utils.reject('Could not create a model.', true));
+        }, data.state).catch(utils.reject('Could not create a model.', true));
     };
 
     /**

jupyter-js-widgets/src/utils.ts

@@ -1,6 +1,8 @@
 // Copyright (c) Jupyter Development Team.
 // Distributed under the terms of the Modified BSD License.
 
+import * as _ from 'underscore';
+
 // TODO: ATTEMPT TO KILL THIS MODULE USING THIRD PARTY LIBRARIES WHEN IPYWIDGETS
 // IS CONVERTED TO NODE COMMONJS.
 
@@ -169,3 +171,111 @@ function escape_html(text: string): string {
     esc.textContent = text;
     return esc.innerHTML;
 };
+
+/**
+ * Takes an object 'state' and fills in buffer[i] at 'path' buffer_paths[i]
+ * where buffer_paths[i] is a list indicating where in the object buffer[i] should
+ * be placed
+ * Example: state = {a: 1, b: {}, c: [0, null]}
+ * buffers = [array1, array2]
+ * buffer_paths = [['b', 'data'], ['c', 1]]
+ * Will lead to {a: 1, b: {data: array1}, c: [0, array2]}
+ */
+export
+function put_buffers(state, buffer_paths, buffers) {
+    for (let i=0; i<buffer_paths.length; i++) {
+        let buffer_path = buffer_paths[i];
+         // say we want to set state[x][y][z] = buffers[i]
+        let obj = state;
+        // we first get obj = state[x][y]
+        for (let j = 0; j < buffer_path.length-1; j++)
+            obj = obj[buffer_path[j]];
+        // and then set: obj[z] = buffers[i]
+        obj[buffer_path[buffer_path.length-1]] = buffers[i];
+    }
+}
+
+/**
+ * The inverse of put_buffers, return an objects with the new state where all buffers(ArrayBuffer)
+ * are removed. If a buffer is a member of an object, that object is cloned, and the key removed. If a buffer
+ * is an element of an array, that array is cloned, and the element is set to null.
+ * See put_buffers for the meaning of buffer_paths
+ * Returns an object with the new state (.state) an array with paths to the buffers (.buffer_paths),
+ * and the buffers associated to those paths (.buffers).
+ */
+export
+function remove_buffers(state) {
+    let buffers = [];
+    let buffer_paths = [];
+    // if we need to remove an object from a list, we need to clone that list, otherwise we may modify
+    // the internal state of the widget model
+    // however, we do not want to clone everything, for performance
+    function remove(obj, path) {
+        if (obj.toJSON) {
+            // We need to get the JSON form of the object before recursing.
+            // See https://developer.mozilla.org/en-US/docs/Web/JavaScript/Reference/Global_Objects/JSON/stringify#toJSON()_behavior
+            obj = obj.toJSON();
+        }
+        if (Array.isArray(obj)) {
+            let is_cloned = false;
+            for (let i = 0; i < obj.length; i++) {
+                let value = obj[i];
+                if(value) {
+                    if (value.buffer instanceof ArrayBuffer || value instanceof ArrayBuffer) {
+                        if(!is_cloned) {
+                            obj = _.clone(obj);
+                            is_cloned = true;
+                        }
+                        buffers.push(value);
+                        buffer_paths.push(path.concat([i]));
+                        // easier to just keep the array, but clear the entry, otherwise we have to think
+                        // about array length, much easier this way
+                        obj[i] = null;
+                    } else {
+                        let new_value  = remove(value, path.concat([i]));
+                        // only assigned when the value changes, we may serialize objects that don't support assignment
+                        if(new_value !== value) {
+                            if(!is_cloned) {
+                                obj = _.clone(obj);
+                                is_cloned = true;
+                            }
+                            obj[i] = new_value;
+                        }
+                    }
+                }
+            }
+        } else if(_.isObject(obj)) {
+            for (let key in obj) {
+                let is_cloned = false;
+                if (obj.hasOwnProperty(key)) {
+                    let value = obj[key];
+                    if(value) {
+                        if (value.buffer instanceof ArrayBuffer || value instanceof ArrayBuffer) {
+                            if(!is_cloned) {
+                                obj = _.clone(obj);
+                                is_cloned = true;
+                            }
+                            buffers.push(value);
+                            buffer_paths.push(path.concat([key]));
+                            delete obj[key]; // for objects/dicts we just delete them
+                        }
+                        else {
+                            let new_value  = remove(value, path.concat([key]));
+                            // only assigned when the value changes, we may serialize objects that don't support assignment
+                            if(new_value !== value) {
+                                if(!is_cloned) {
+                                    obj = _.clone(obj);
+                                    is_cloned = true;
+                                }
+                                obj[key] = new_value;
+                            }
+                        }
+                    }
+                }
+            }
+        }
+        return obj;
+    }
+    let new_state = remove(state, []);
+    return {state: new_state, buffers: buffers, buffer_paths: buffer_paths}
+}

jupyter-js-widgets/src/widget.ts

@@ -168,12 +168,11 @@ class WidgetModel extends Backbone.Model {
             case 'update':
                 this.state_change = this.state_change
                     .then(() => {
-                        var state = msg.content.data.state || {};
-                        var buffer_keys = msg.content.data.buffers || [];
+                        // see Widget.open/_split_state_buffers about why we need state_with_buffers
+                        var state = msg.content.data.state;
+                        var buffer_paths = msg.content.data.buffer_paths || [];
                         var buffers = msg.buffers || [];
-                        for (var i=0; i<buffer_keys.length; i++) {
-                            state[buffer_keys[i]] = buffers[i];
-                        }
+                        utils.put_buffers(state, buffer_paths, buffers);
                         return (this.constructor as typeof WidgetModel)._deserialize_state(state, this.widget_manager);
                     }).then((state) => {
                         this.set_state(state);
@@ -394,25 +393,15 @@ class WidgetModel extends Backbone.Model {
         utils.resolvePromisesDict(attrs).then((state) => {
             // get binary values, then send
             var keys = Object.keys(state);
-            var buffers = [];
-            var buffer_keys = [];
-            for (var i=0; i<keys.length; i++) {
-                var key = keys[i];
-                var value = state[key];
-                if (value) {
-                    if (value.buffer instanceof ArrayBuffer
-                        || value instanceof ArrayBuffer) {
-                        buffers.push(value);
-                        buffer_keys.push(key);
-                        delete state[key];
-                    }
-                }
-            }
+            // this function goes through lists and object and removes arraybuffers
+            // they will be transferred separately, since they don't json'ify
+            // on the python side the inverse happens
+            var split = utils.remove_buffers(state);
             this.comm.send({
                 method: 'backbone',
-                sync_data: state,
-                buffer_keys: buffer_keys
-            }, callbacks, {}, buffers);
+                sync_data: split.state,
+                buffer_paths: split.buffer_paths
+            }, callbacks, {}, split.buffers);
         }).catch((error) => {
             this.pending_msgs--;
             return (utils.reject('Could not send widget sync message', true))(error);