Merge pull request #2022 from jasongrout/callbackid-4.x

Backport #1958 to 4.x: Add an output callback override stack

Author: SylvainCorlay

notebook/static/services/kernels/comm.js

@@ -129,12 +129,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];
     };
@@ -194,15 +191,15 @@ 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);
             }
         }
     };
 
     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

@@ -42,6 +42,7 @@ define([
         this.username = "username";
         this.session_id = utils.uuid();
         this._msg_callbacks = {};
+        this._msg_callbacks_overrides = {};
         this._msg_queue = Promise.resolve();
         this.info_reply = {}; // kernel_info_reply stored here after starting
 
@@ -301,6 +302,8 @@ define([
         this.events.trigger('kernel_restarting.Kernel', {kernel: this});
         this.stop_channels();
 
+        this._msg_callbacks = {};
+        this._msg_callbacks_overrides = {};
         var that = this;
         var on_success = function (data, status, xhr) {
             that.events.trigger('kernel_created.Kernel', {kernel: that});
@@ -847,6 +850,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.
      *
@@ -891,10 +923,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 = {};
@@ -903,11 +941,31 @@ define([
             cbcopy.input = callbacks.input;
             cbcopy.shell_done = (!callbacks.shell);
             cbcopy.iopub_done = (!callbacks.iopub);
-        } 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() {
@@ -1032,7 +1090,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;
         }
@@ -1048,7 +1106,8 @@ define([
      * @function _handle_output_message
      */
     Kernel.prototype._handle_output_message = function (msg) {
-        var callbacks = this.get_callbacks_for_msg(msg.parent_header.msg_id);
+        var msg_id = msg.parent_header.msg_id;
+        var callbacks = this.get_output_callbacks_for_msg(msg_id);
         if (!callbacks || !callbacks.iopub) {
             // The message came from another client. Let the UI decide what to
             // do with it.

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);
+    });
 });