-
-
Notifications
You must be signed in to change notification settings - Fork 4.4k
/
routers.py
58 lines (43 loc) · 1.68 KB
/
routers.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
import uuid
from Queue import Queue
import json
class ZMQStreamRouter(object):
def __init__(self, zmq_stream, session):
self.zmq_stream = zmq_stream
self.session = session
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:
print "Got message: ", msg
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, session):
ZMQStreamRouter.__init__(self, zmq_stream, session)
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):
print "Inbound message: ", msg
self._request_queue.put(client_id)
self.session.send(self.zmq_stream, msg)