IPython/frontend/html/notebook/routers.py

import uuid
from Queue import Queue


class ZMQStreamRouter(object):

    def __init__(self, zmq_stream):
        self.zmq_stream = zmq_stream
        self._clients = {}
        self.zmq_stream.on_recv(self._on_zmq_reply)

    def register_client(self, client):
        client_id = uuid.uuid4()
        self._clients[client_id] = client
        return client_id

    def unregister_client(self, client_id):
        del self._clients[client_id]

    def copy_clients(self, router):
        # Copy the clients of another router.
        for client_id, client in router._clients.items():
            client.router = self
            self._clients[client_id] = client


class IOPubStreamRouter(ZMQStreamRouter):

    def _on_zmq_reply(self, msg_list):
        for client_id, client in self._clients.items():
            for msg in msg_list:
                client.write_message(msg)

    def forward_unicode(self, client_id, msg):
        # This is a SUB stream that we should never write to.
        pass


class ShellStreamRouter(ZMQStreamRouter):

    def __init__(self, zmq_stream):
        ZMQStreamRouter.__init__(self, zmq_stream)
        self._request_queue = Queue()

    def _on_zmq_reply(self, msg_list):
        client_id = self._request_queue.get(block=False)
        client = self._clients.get(client_id)
        if client is not None:
            for msg in msg_list:
                client.write_message(msg)

    def forward_unicode(self, client_id, msg):
        self._request_queue.put(client_id)
        self.zmq_stream.send_unicode(msg)