Merge pull request #1958 from jasongrout/callbackid

Add output callback overrides

notebook/static/services/kernels/comm.js

@@ -128,12 +128,9 @@ define([
         }
 
         this.comms[content.comm_id] = this.comms[content.comm_id].then(function(comm) {
-            try {
-                comm.handle_msg(msg);
-            } catch (e) {
-                console.log("Exception handling comm msg: ", e, e.stack, msg);
-            }
-            return comm;
+            return (Promise.resolve(comm.handle_msg(msg))
+                .catch(utils.reject('Exception handling comm message'))
+                .then(function() {return comm;}));
         });
         return this.comms[content.comm_id];
     };
@@ -193,7 +190,7 @@ define([
         var callback = this['_' + key + '_callback'];
         if (callback) {
             try {
-                callback(msg);
+                return callback(msg);
             } catch (e) {
                 console.log("Exception in Comm callback", e, e.stack, msg);
             }
@@ -201,7 +198,7 @@ define([
     };
     
     Comm.prototype.handle_msg = function (msg) {
-        this._callback('msg', msg);
+        return this._callback('msg', msg);
     };
     
     Comm.prototype.handle_close = function (msg) {

notebook/static/services/kernels/kernel.js

@@ -41,6 +41,7 @@ define([
         this.username = "username";
         this.session_id = utils.uuid();
         this._msg_callbacks = {};
+        this._msg_callbacks_overrides = {};
         this._display_id_to_parent_ids = {};
         this._msg_queue = Promise.resolve();
         this.info_reply = {}; // kernel_info_reply stored here after starting
@@ -301,6 +302,7 @@ define([
         this.events.trigger('kernel_restarting.Kernel', {kernel: this});
         this.stop_channels();
         this._msg_callbacks = {};
+        this._msg_callbacks_overrides = {};
         this._display_id_to_parent_ids = {};
 
         var that = this;
@@ -850,6 +852,35 @@ define([
         }
     };
 
+    /**
+     * Get output callbacks for a specific message.
+     *
+     * @function get_output_callbacks_for_msg
+     *
+     * Since output callbacks can be overridden, we first check the override stack.
+     */
+    Kernel.prototype.get_output_callbacks_for_msg = function (msg_id) {
+        return this.get_callbacks_for_msg(this.get_output_callback_id(msg_id));
+    };
+
+
+    /**
+     * Get the output callback id for a message
+     *
+     * Since output callbacks can be redirected, this may not be the same as
+     * the msg_id.
+     *
+     * @function get_output_callback_id
+     */
+    Kernel.prototype.get_output_callback_id = function (msg_id) {
+        var callback_id = msg_id;
+        var overrides = this._msg_callbacks_overrides[msg_id];
+        if (overrides && overrides.length > 0) {
+            callback_id = overrides[overrides.length-1];
+        }
+        return callback_id
+    }
+
     /**
      * Clear callbacks for a specific message.
      *
@@ -913,10 +944,16 @@ define([
      * 
      * }
      *
+     * If the third parameter is truthy, the callback is set as the last
+     * callback registered.
+     *
      * @function set_callbacks_for_msg
      */
-    Kernel.prototype.set_callbacks_for_msg = function (msg_id, callbacks) {
-        this.last_msg_id = msg_id;
+    Kernel.prototype.set_callbacks_for_msg = function (msg_id, callbacks, save) {
+        var remember = save || true;
+        if (remember) {
+            this.last_msg_id = msg_id;
+        }
         if (callbacks) {
             // shallow-copy mapping, because we will modify it at the top level
             var cbcopy = this._msg_callbacks[msg_id] = this.last_msg_callbacks = {};
@@ -931,11 +968,31 @@ define([
                 // default to clear-on-done
                 cbcopy.clear_on_done = true;
             }
-        } else {
+        } else if (remember) {
             this.last_msg_callbacks = {};
         }
     };
-    
+
+    /**
+     * Override output callbacks for a particular msg_id
+     */
+    Kernel.prototype.output_callback_overrides_push = function(msg_id, callback_id) {
+        var output_callbacks = this._msg_callbacks_overrides[msg_id];
+        if (output_callbacks === void 0) {
+            this._msg_callbacks_overrides[msg_id] = output_callbacks = [];
+        }
+        output_callbacks.push(callback_id);
+    }
+
+    Kernel.prototype.output_callback_overrides_pop = function(msg_id) {
+        var callback_ids = this._msg_callbacks_overrides[msg_id];
+        if (!callback_ids) {
+            console.error("Popping callback overrides, but none registered", msg_id);
+            return;
+        }
+        return callback_ids.pop();
+    }
+
     Kernel.prototype._handle_ws_message = function (e) {
         var that = this;
         this._msg_queue = this._msg_queue.then(function() {
@@ -1060,7 +1117,7 @@ define([
      * @function _handle_clear_output
      */
     Kernel.prototype._handle_clear_output = function (msg) {
-        var callbacks = this.get_callbacks_for_msg(msg.parent_header.msg_id);
+        var callbacks = this.get_output_callbacks_for_msg(msg.parent_header.msg_id);
         if (!callbacks || !callbacks.iopub) {
             return;
         }
@@ -1078,7 +1135,7 @@ define([
     Kernel.prototype._handle_output_message = function (msg) {
         var that = this;
         var msg_id = msg.parent_header.msg_id;
-        var callbacks = this.get_callbacks_for_msg(msg_id);
+        var callbacks = this.get_output_callbacks_for_msg(msg_id);
         if (['display_data', 'update_display_data', 'execute_result'].indexOf(msg.header.msg_type) > -1) {
             // display_data messages may re-route based on their display_id
             var display_id = (msg.content.transient || {}).display_id;
@@ -1110,8 +1167,9 @@ define([
                 if (this._display_id_to_parent_ids[display_id] === undefined) {
                     this._display_id_to_parent_ids[display_id] = [];
                 }
-                if (this._display_id_to_parent_ids[display_id].indexOf(msg_id) === -1) {
-                    this._display_id_to_parent_ids[display_id].push(msg_id);
+                var callback_id = this.get_output_callback_id(msg_id);
+                if (this._display_id_to_parent_ids[display_id].indexOf(callback_id) === -1) {
+                    this._display_id_to_parent_ids[display_id].push(callback_id);
                 }
                 // and in callbacks for cleanup on clear_callbacks_for_msg
                 if (callbacks && callbacks.display_ids.indexOf(display_id) === -1) {

notebook/tests/notebook/display_id.js

@@ -79,4 +79,89 @@ casper.notebook_test(function () {
         this.test.assertEquals(outputs2[2].data['text/plain'], '4', 'output[2][2]');
     });
 
+    this.then(function () {
+        this.echo("Test output callback overrides work with display ids");
+    });
+
+    this.thenEvaluate(function () {
+        Jupyter.notebook.insert_cell_at_index("code", 3);
+        var cell = Jupyter.notebook.get_cell(3);
+        cell.set_text([
+            "display_with_id(5, 'here')",
+            "display_with_id(6, 'here', update=True)",
+        ].join('\n'));
+        cell.execute();
+        var kernel = IPython.notebook.kernel;
+        var msg_id = cell.last_msg_id;
+        var callback_id = 'mycallbackid'
+        cell.iopub_messages = [];
+        var add_msg = function(msg) {
+            msg.content.output_type = msg.msg_type;
+            cell.iopub_messages.push(msg.content);
+        };
+        kernel.set_callbacks_for_msg(callback_id, {
+            iopub: {
+                output: add_msg,
+                clear_output: add_msg,
+            }
+        }, false);
+        kernel.output_callback_overrides_push(msg_id, callback_id);
+    });
+
+    this.wait_for_idle();
+
+    this.then(function () {
+        var returned = this.evaluate(function () {
+            var cell = IPython.notebook.get_cell(3);
+            return [cell.output_area.outputs, cell.iopub_messages];
+        });
+        var cell_results = returned[0];
+        var callback_results = returned[1];
+        this.test.assertEquals(cell_results.length, 0, "correct number of cell outputs");
+        this.test.assertEquals(callback_results.length, 2, "correct number of callback outputs");
+        this.test.assertEquals(callback_results[0].output_type, 'display_data', 'check output_type 0');
+        this.test.assertEquals(callback_results[0].transient.display_id, 'here', 'check display id 0');
+        this.test.assertEquals(callback_results[0].data['text/plain'], '5', 'value');
+        this.test.assertEquals(callback_results[1].output_type, 'update_display_data', 'check output_type 1');
+        this.test.assertEquals(callback_results[1].transient.display_id, 'here', 'display id 1');
+        this.test.assertEquals(callback_results[1].data['text/plain'], '6', 'value');
+    });
+
+    this.thenEvaluate(function () {
+        Jupyter.notebook.insert_cell_at_index("code", 4);
+        var cell = Jupyter.notebook.get_cell(4);
+        cell.set_text([
+            "display_with_id(7, 'here')",
+            "display_with_id(8, 'here', update=True)",
+        ].join('\n'));
+        cell.execute();
+    });
+
+    this.wait_for_idle();
+
+    this.then(function () {
+        var returned = JSON.parse(this.evaluate(function () {
+            var cell3 = Jupyter.notebook.get_cell(3);
+            var cell4 = Jupyter.notebook.get_cell(4);
+            return JSON.stringify([cell4.output_area.outputs, cell3.iopub_messages]);
+        }));
+        var cell_results = returned[0];
+        var callback_results = returned[1];
+        this.test.assertEquals(cell_results.length, 1, "correct number of cell outputs");
+        this.test.assertEquals(callback_results.length, 4, "correct number of callback outputs");
+        this.test.assertEquals(callback_results[0].output_type, 'display_data', 'check output_type 0');
+        this.test.assertEquals(callback_results[0].transient.display_id, 'here', 'check display id 0');
+        this.test.assertEquals(callback_results[0].data['text/plain'], '5', 'value');
+        this.test.assertEquals(callback_results[1].output_type, 'update_display_data', 'check output_type 1');
+        this.test.assertEquals(callback_results[1].transient.display_id, 'here', 'display id 1');
+        this.test.assertEquals(callback_results[1].data['text/plain'], '6', 'value');
+        this.test.assertEquals(callback_results[2].output_type, 'display_data', 'check output_type 2');
+        this.test.assertEquals(callback_results[2].transient.display_id, 'here', 'check display id 2');
+        this.test.assertEquals(callback_results[2].data['text/plain'], '7', 'value');
+        this.test.assertEquals(callback_results[3].output_type, 'update_display_data', 'check output_type 3');
+        this.test.assertEquals(callback_results[3].transient.display_id, 'here', 'display id 3');
+        this.test.assertEquals(callback_results[3].data['text/plain'], '8', 'value');
+    });
+
+
 });

notebook/tests/notebook/output.js

@@ -4,6 +4,17 @@
 
 casper.notebook_test(function () {
     
+    this.compare_outputs = function(results, expected) {
+        for (var i = 0; i < results.length; i++) {
+            var r = results[i];
+            var ex = expected[i];
+            this.test.assertEquals(r.output_type, ex.output_type, "output  " + i + " = " + r.output_type);
+            if (r.output_type === 'stream') {
+                this.test.assertEquals(r.name, ex.name, "stream  " + i + " = " + r.name);
+                this.test.assertEquals(r.text, ex.text, "content " + i);
+            }
+        }
+    }
     this.test_coalesced_output = function (msg, code, expected) {
         this.then(function () {
             this.echo("Test coalesced output: " + msg);
@@ -24,31 +35,23 @@ casper.notebook_test(function () {
                 return cell.output_area.outputs;
             });
             this.test.assertEquals(results.length, expected.length, "correct number of outputs");
-            for (var i = 0; i < results.length; i++) {
-                var r = results[i];
-                var ex = expected[i];
-                this.test.assertEquals(r.output_type, ex.output_type, "output  " + i);
-                if (r.output_type === 'stream') {
-                    this.test.assertEquals(r.name, ex.name, "stream  " + i);
-                    this.test.assertEquals(r.text, ex.text, "content " + i);
-                }
-            }
+            this.compare_outputs(results, expected);
         });
         
     };
-    
+
     this.thenEvaluate(function () {
         IPython.notebook.insert_cell_at_index("code", 0);
         var cell = IPython.notebook.get_cell(0);
         cell.set_text([
             "from __future__ import print_function",
             "import sys",
-            "from IPython.display import display"
+            "from IPython.display import display, clear_output"
             ].join("\n")
         );
         cell.execute();
     });
-    
+
     this.test_coalesced_output("stdout", [
         "print(1)",
         "sys.stdout.flush()",
@@ -123,4 +126,134 @@ casper.notebook_test(function () {
             },
         }]
     );
+
+    this.then(function () {
+        this.echo("Test output callback overrides");
+    });
+
+    this.thenEvaluate(function () {
+        IPython.notebook.insert_cell_at_index("code", 0);
+        var cell = IPython.notebook.get_cell(0);
+        cell.set_text(["print(1)",
+        "sys.stdout.flush()",
+        "print(2)",
+        "sys.stdout.flush()",
+        "print(3, file=sys.stderr)",
+        "sys.stdout.flush()",
+        "display(2)",
+        "clear_output()",
+        "sys.stdout.flush()",
+        "print('remove handler')",
+        "sys.stdout.flush()",
+        "print('back to cell')",
+        "sys.stdout.flush()",
+        ].join('\n'));
+        cell.execute();
+        var kernel = IPython.notebook.kernel;
+        var msg_id = cell.last_msg_id;
+        var callback_id = 'mycallbackid'
+        cell.iopub_messages = [];
+        var add_msg = function(msg) {
+            if (msg.content.text==="remove handler\n") {
+                kernel.output_callback_overrides_pop(msg_id);
+            }
+            msg.content.output_type = msg.msg_type;
+            cell.iopub_messages.push(msg.content);
+        };
+        kernel.set_callbacks_for_msg(callback_id, {
+            iopub: {
+                output: add_msg,
+                clear_output: add_msg,
+            }
+        }, false);
+        kernel.output_callback_overrides_push(msg_id, callback_id);
+    });
+    
+    this.wait_for_idle();
+    
+    this.then(function () {
+        var expected_callback = [{
+            output_type: "stream",
+            name: "stdout",
+            text: "1\n"
+        }, {
+            output_type: "stream",
+            name: "stdout",
+            text: "2\n"
+        }, {
+            output_type: "stream",
+            name: "stderr",
+            text: "3\n"
+        },{
+            output_type: "display_data",
+        },{
+            output_type: "clear_output",
+        },{
+            output_type: "stream",
+            name: "stdout",
+            text: "remove handler\n"
+        },]
+        var expected_cell = [{
+            output_type: "stream",
+            name: "stdout",
+            text: "back to cell\n"
+        }]
+        var returned = this.evaluate(function () {
+            var cell = IPython.notebook.get_cell(0);
+            return [cell.output_area.outputs, cell.iopub_messages];
+        });
+        var cell_results = returned[0];
+        var callback_results = returned[1];
+        this.test.assertEquals(cell_results.length, expected_cell.length, "correct number of cell outputs");
+        this.test.assertEquals(callback_results.length, expected_callback.length, "correct number of callback outputs");
+        this.compare_outputs(cell_results, expected_cell);
+        this.compare_outputs(callback_results, expected_callback);
+    });
+
+    this.then(function () {
+        this.echo("Test output callback overrides get execute_results messages too");
+    });
+
+    this.thenEvaluate(function () {
+        IPython.notebook.insert_cell_at_index("code", 0);
+        var cell = IPython.notebook.get_cell(0);
+        cell.set_text("'end'");
+        cell.execute();
+        var kernel = IPython.notebook.kernel;
+        var msg_id = cell.last_msg_id;
+        var callback_id = 'mycallbackid2'
+        cell.iopub_messages = [];
+        var add_msg = function(msg) {
+            msg.content.output_type = msg.msg_type;
+            cell.iopub_messages.push(msg.content);
+        };
+        kernel.set_callbacks_for_msg(callback_id, {
+            iopub: {
+                output: add_msg,
+                clear_output: add_msg,
+            }
+        }, false);
+        kernel.output_callback_overrides_push(msg_id, callback_id);
+    });
+    
+    this.wait_for_idle();
+    
+    this.then(function () {
+        var expected_callback = [{
+            output_type: "execute_result",
+            data: {
+            "text/plain" : "'end'"
+            }
+        }];
+        var expected_cell = [];
+        var returned = this.evaluate(function () {
+            var cell = IPython.notebook.get_cell(0);
+            return [cell.output_area.outputs, cell.iopub_messages];
+        });
+        var cell_results = returned[0];
+        var callback_results = returned[1];
+        this.test.assertEquals(cell_results.length, expected_cell.length, "correct number of cell outputs");
+        this.test.assertEquals(callback_results.length, expected_callback.length, "correct number of callback outputs");
+        this.compare_outputs(callback_results, expected_callback);
+    });
 });