ipython · ellisonbg · Apr 26, 2013 · Mar 19, 2013 · Mar 19, 2013 · Mar 19, 2013
diff --git a/IPython/frontend/html/notebook/handlers.py b/IPython/frontend/html/notebook/handlers.py
@@ -506,6 +506,12 @@ def _inject_cookie_message(self, msg):
             # Cookie constructor doesn't accept unicode strings
             # under Python 2.x for some reason
             msg = msg.encode('utf8', 'replace')
+        try:
+            identity, msg = msg.split(':', 1)
+            self.session.session = identity.decode('ascii')
+        except Exception:
+            logging.error("First ws message didn't have the form 'identity:[cookie]' - %r", msg)
+
         try:
             self.request._cookies = Cookie.SimpleCookie(msg)
         except:
@@ -519,53 +525,61 @@ def on_first_message(self, msg):
         self.on_message = self.save_on_message
 
 
-class IOPubHandler(AuthenticatedZMQStreamHandler):
-
+class ZMQChannelHandler(AuthenticatedZMQStreamHandler):
+
+    @property
+    def max_msg_size(self):
+        return self.settings.get('max_msg_size', 65535)
+
+    def create_stream(self):
+        km = self.kernel_manager
+        meth = getattr(km, 'connect_%s' % self.channel)
+        self.zmq_stream = meth(self.kernel_id, identity=self.session.bsession)
+
     def initialize(self, *args, **kwargs):
-        self.iopub_stream = None
-
+        self.zmq_stream = None
+    
     def on_first_message(self, msg):
         try:
-            super(IOPubHandler, self).on_first_message(msg)
+            super(ZMQChannelHandler, self).on_first_message(msg)
         except web.HTTPError:
             self.close()
             return
-        km = self.kernel_manager
-        kernel_id = self.kernel_id
-        km.add_restart_callback(kernel_id, self.on_kernel_restarted)
-        km.add_restart_callback(kernel_id, self.on_restart_failed, 'dead')
         try:
-            self.iopub_stream = km.connect_iopub(kernel_id)
+            self.create_stream()
         except web.HTTPError:
             # WebSockets don't response to traditional error codes so we
             # close the connection.
             if not self.stream.closed():
                 self.stream.close()
             self.close()
         else:
-            self.iopub_stream.on_recv(self._on_zmq_reply)
+            self.zmq_stream.on_recv(self._on_zmq_reply)
 
     def on_message(self, msg):
-        pass
-
-    def _send_status_message(self, status):
-        msg = self.session.msg("status",
-            {'execution_state': status}
-        )
-        self.write_message(jsonapi.dumps(msg, default=date_default))
-
-    def on_kernel_restarted(self):
-        self.log.warn("kernel %s restarted", self.kernel_id)
-        self._send_status_message('restarting')
-
-    def on_restart_failed(self):
-        self.log.error("kernel %s restarted failed!", self.kernel_id)
-        self._send_status_message('dead')
+        if len(msg) < self.max_msg_size:
+            msg = jsonapi.loads(msg)
+            self.session.send(self.zmq_stream, msg)
 
        def on_close(self):
         # This method can be called twice, once by self.kernel_died and once
         # from the WebSocket close event. If the WebSocket connection is
         # closed before the ZMQ streams are setup, they could be None.
+        if self.zmq_stream is not None and not self.zmq_stream.closed():
+            self.zmq_stream.on_recv(None)
+            self.zmq_stream.close()
+
+
+class IOPubHandler(ZMQChannelHandler):
+    channel = 'iopub'
+
+    def create_stream(self):
+        super(IOPubHandler, self).create_stream()
+        km = self.kernel_manager
+        km.add_restart_callback(self.kernel_id, self.on_kernel_restarted)
+        km.add_restart_callback(self.kernel_id, self.on_restart_failed, 'dead')
+
+    def on_close(self):
         km = self.kernel_manager
         if self.kernel_id in km:
             km.remove_restart_callback(
@@ -574,48 +588,31 @@ def on_close(self):
             km.remove_restart_callback(
                 self.kernel_id, self.on_restart_failed, 'dead',
             )
-        if self.iopub_stream is not None and not self.iopub_stream.closed():
-            self.iopub_stream.on_recv(None)
-            self.iopub_stream.close()
-
-
-class ShellHandler(AuthenticatedZMQStreamHandler):
+        super(IOPubHandler, self).on_close()
 
-    @property
-    def max_msg_size(self):
-        return self.settings.get('max_msg_size', 65535)
-
-    def initialize(self, *args, **kwargs):
-        self.shell_stream = None
+    def _send_status_message(self, status):
+        msg = self.session.msg("status",
+            {'execution_state': status}
+        )
+        self.write_message(jsonapi.dumps(msg, default=date_default))
 
-    def on_first_message(self, msg):
-        try:
-            super(ShellHandler, self).on_first_message(msg)
-        except web.HTTPError:
-            self.close()
-            return
-        km = self.kernel_manager
-        kernel_id = self.kernel_id
-        try:
-            self.shell_stream = km.connect_shell(kernel_id)
-        except web.HTTPError:
-            # WebSockets don't response to traditional error codes so we
-            # close the connection.
-            if not self.stream.closed():
-                self.stream.close()
-            self.close()
-        else:
-            self.shell_stream.on_recv(self._on_zmq_reply)
+    def on_kernel_restarted(self):
+        logging.warn("kernel %s restarted", self.kernel_id)
+        self._send_status_message('restarting')
 
+    def on_restart_failed(self):
+        logging.error("kernel %s restarted failed!", self.kernel_id)
+        self._send_status_message('dead')
+
     def on_message(self, msg):
-        if len(msg) < self.max_msg_size:
-            msg = jsonapi.loads(msg)
-            self.session.send(self.shell_stream, msg)
+        """IOPub messages make no sense"""
+        pass
 
-    def on_close(self):
-        # Make sure the stream exists and is not already closed.
-        if self.shell_stream is not None and not self.shell_stream.closed():
-            self.shell_stream.close()
+class ShellHandler(ZMQChannelHandler):
+    channel = 'shell'
+
+class StdinHandler(ZMQChannelHandler):
+    channel = 'stdin'
 
 
 #-----------------------------------------------------------------------------

diff --git a/IPython/frontend/html/notebook/notebookapp.py b/IPython/frontend/html/notebook/notebookapp.py
@@ -66,7 +66,7 @@
 from .kernelmanager import MappingKernelManager
 from .handlers import (LoginHandler, LogoutHandler,
     ProjectDashboardHandler, NewHandler, NamedNotebookHandler,
-    MainKernelHandler, KernelHandler, KernelActionHandler, IOPubHandler,
+    MainKernelHandler, KernelHandler, KernelActionHandler, IOPubHandler, StdinHandler,
     ShellHandler, NotebookRootHandler, NotebookHandler, NotebookCopyHandler,
     RSTHandler, AuthenticatedFileHandler, PrintNotebookHandler,
     MainClusterHandler, ClusterProfileHandler, ClusterActionHandler,
@@ -160,6 +160,7 @@ def __init__(self, ipython_app, kernel_manager, notebook_manager,
             (r"/kernels/%s/%s" % (_kernel_id_regex, _kernel_action_regex), KernelActionHandler),
             (r"/kernels/%s/iopub" % _kernel_id_regex, IOPubHandler),
             (r"/kernels/%s/shell" % _kernel_id_regex, ShellHandler),
+            (r"/kernels/%s/stdin" % _kernel_id_regex, StdinHandler),
             (r"/notebooks", NotebookRootHandler),
             (r"/notebooks/%s" % _notebook_id_regex, NotebookHandler),
             (r"/rstservice/render", RSTHandler),

diff --git a/IPython/frontend/html/notebook/static/css/style.min.css b/IPython/frontend/html/notebook/static/css/style.min.css
diff --git a/IPython/frontend/html/notebook/static/js/codecell.js b/IPython/frontend/html/notebook/static/js/codecell.js
@@ -245,7 +245,8 @@ var IPython = (function (IPython) {
             'execute_reply': $.proxy(this._handle_execute_reply, this),
             'output': $.proxy(this.output_area.handle_output, this.output_area),
             'clear_output': $.proxy(this.output_area.handle_clear_output, this.output_area),
-            'set_next_input': $.proxy(this._handle_set_next_input, this)
+            'set_next_input': $.proxy(this._handle_set_next_input, this),
+            'input_request': $.proxy(this._handle_input_request, this)
         };
         var msg_id = this.kernel.execute(this.get_text(), callbacks, {silent: false});
     };
@@ -260,10 +261,23 @@ var IPython = (function (IPython) {
         $([IPython.events]).trigger('set_dirty.Notebook', {'value': true});
     }
 
+    /**
+     * @method _handle_set_next_input
+     * @private
+     */
     CodeCell.prototype._handle_set_next_input = function (text) {
         var data = {'cell': this, 'text': text}
         $([IPython.events]).trigger('set_next_input.Notebook', data);
     }
+
+    /**
+     * @method _handle_input_request
+     * @private
+     */
+    CodeCell.prototype._handle_input_request = function (content) {
+        this.output_area.append_raw_input(content);
+    }
+
 
     // Basic cell manipulation.
 

diff --git a/IPython/frontend/html/notebook/static/js/kernel.js b/IPython/frontend/html/notebook/static/js/kernel.js
@@ -28,6 +28,7 @@ var IPython = (function (IPython) {
         this.kernel_id = null;
         this.shell_channel = null;
         this.iopub_channel = null;
+        this.stdin_channel = null;
         this.base_url = base_url;
         this.running = false;
         this.username = "username";
@@ -127,9 +128,12 @@ var IPython = (function (IPython) {
         var ws_url = this.ws_url + this.kernel_url;
         console.log("Starting WebSockets:", ws_url);
         this.shell_channel = new this.WebSocket(ws_url + "/shell");
+        this.stdin_channel = new this.WebSocket(ws_url + "/stdin");
         this.iopub_channel = new this.WebSocket(ws_url + "/iopub");
         send_cookie = function(){
-            this.send(document.cookie);
+            // send the session id so the Session object Python-side
+            // has the same identity
+            this.send(that.session_id + ':' + document.cookie);
         };
         var already_called_onclose = false; // only alert once
         var ws_closed_early = function(evt){
@@ -150,38 +154,41 @@ var IPython = (function (IPython) {
                 that._websocket_closed(ws_url, false);
             }
         };
-        this.shell_channel.onopen = send_cookie;
-        this.shell_channel.onclose = ws_closed_early;
-        this.iopub_channel.onopen = send_cookie;
-        this.iopub_channel.onclose = ws_closed_early;
+        var channels = [this.shell_channel, this.iopub_channel, this.stdin_channel];
+        for (var i=0; i < channels.length; i++) {
+            channels[i].onopen = send_cookie;
+            channels[i].onclose = ws_closed_early;
+        }
         // switch from early-close to late-close message after 1s
         setTimeout(function() {
-            if (that.shell_channel !== null) {
-                that.shell_channel.onclose = ws_closed_late;
-            }
-            if (that.iopub_channel !== null) {
-                that.iopub_channel.onclose = ws_closed_late;
+            for (var i=0; i < channels.length; i++) {
+                if (channels[i] !== null) {
+                    channels[i].onclose = ws_closed_late;
+                }
             }
         }, 1000);
-        this.shell_channel.onmessage = $.proxy(this._handle_shell_reply,this);
-        this.iopub_channel.onmessage = $.proxy(this._handle_iopub_reply,this);
+        this.shell_channel.onmessage = $.proxy(this._handle_shell_reply, this);
+        this.iopub_channel.onmessage = $.proxy(this._handle_iopub_reply, this);
+        this.stdin_channel.onmessage = $.proxy(this._handle_input_request, this);
+
+        $([IPython.events]).on('send_input_reply.Kernel', function(evt, data) { 
+            that.send_input_reply(data);
+        });
     };
 
     /**
      * Start the `shell`and `iopub` channels.
      * @method stop_channels
      */
     Kernel.prototype.stop_channels = function () {
-        if (this.shell_channel !== null) {
-            this.shell_channel.onclose = function (evt) {};
-            this.shell_channel.close();
-            this.shell_channel = null;
-        };
-        if (this.iopub_channel !== null) {
-            this.iopub_channel.onclose = function (evt) {};
-            this.iopub_channel.close();
-            this.iopub_channel = null;
+        var channels = [this.shell_channel, this.iopub_channel, this.stdin_channel];
+        for (var i=0; i < channels.length; i++) {
+            if ( channels[i] !== null ) {
+                channels[i].onclose = function (evt) {};
+                channels[i].close();
+            }
         };
+        this.shell_channel = this.iopub_channel = this.stdin_channel = null;
     };
 
     // Main public methods.
@@ -284,6 +291,9 @@ var IPython = (function (IPython) {
             user_expressions : {},
             allow_stdin : false
         };
+        if (callbacks.input_request !== undefined) {
+            content.allow_stdin = true;
+        }
         $.extend(true, content, options)
         $([IPython.events]).trigger('execution_request.Kernel', {kernel: this, content:content});
         var msg = this._get_msg("execute_request", content);
@@ -343,8 +353,18 @@ var IPython = (function (IPython) {
         };
     };
 
+    Kernel.prototype.send_input_reply = function (input) {
+        var content = {
+            value : input,
+        };
+        $([IPython.events]).trigger('input_reply.Kernel', {kernel: this, content:content});
+        var msg = this._get_msg("input_reply", content);
+        this.stdin_channel.send(JSON.stringify(msg));
+        return msg.header.msg_id;
+    };
+
 
-    // Reply handlers.
+    // Reply handlers
 
     Kernel.prototype.get_callbacks_for_msg = function (msg_id) {
         var callbacks = this._msg_callbacks[msg_id];
@@ -433,6 +453,26 @@ var IPython = (function (IPython) {
     };
 
 
+    Kernel.prototype._handle_input_request = function (e) {
+        var request = $.parseJSON(e.data);
+        var header = request.header;
+        var content = request.content;
+        var metadata = request.metadata;
+        var msg_type = header.msg_type;
+        if (msg_type !== 'input_request') {
+            console.log("Invalid input request!", request);
+            return;
+        }
+        var callbacks = this.get_callbacks_for_msg(request.parent_header.msg_id);
+        if (callbacks !== undefined) {
+            var cb = callbacks[msg_type];
+            if (cb !== undefined) {
+                cb(content, metadata);
+            }
+        };
+    };
+
+
     IPython.Kernel = Kernel;
 
     return IPython;