import zmq
import threading
from hedgehog.protocol import sockets, utils


class HedgehogServer(threading.Thread):
    def __init__(self, endpoint, handler, context=None):
        super(HedgehogServer, self).__init__()
        self.context = context or zmq.Context.instance()
        self.endpoint = endpoint
        self.handler = handler
        self.killer = utils.Killer()

    def close(self):
        self.killer.kill()

    def run(self):
        socket = self.context.socket(zmq.ROUTER)
        socket.bind(self.endpoint)
        socket = sockets.RouterWrapper(socket)
        killer = self.killer.connect_receiver()

        while True:
            pollin, _, _ = zmq.select([socket.socket, killer], [], [])
            if pollin:
                if socket.socket in pollin:
                    ident, msg = socket.recv()
                    try:
                        handler = getattr(self.handler, msg._command_oneof)
                    except AttributeError:
                        # TODO handle unknown commands
                        print(msg._command_oneof + ': unknown command')
                        pass
                    else:
                        handler(socket, ident, msg)
                if killer in pollin:
                    break

        socket.close()
        killer.close()

@@ -365,24 +365,431 @@
 ils.Killer()
+%0A        self._poller = None%0A        self.socket = None%0A%0A    @property%0A    def poller(self):%0A        return self._poller%5B0%5D%0A%0A    @property%0A    def sockets(self):%0A        return self._poller%5B1%5D%0A%0A    def register(self, socket, cb):%0A        self.poller.register(socket)%0A        self.sockets%5Bsocket%5D = cb%0A%0A    def unregister(self, socket):%0A        self.poller.unregister(socket)%0A        del self.sockets%5Bsocket%5D
 %0A%0A    def cl
@@ -838,32 +838,73 @@
  def run(self):%0A
+        self._poller = zmq.Poller(), %7B%7D%0A%0A
         socket =
@@ -971,32 +971,37 @@
 dpoint)%0A        
+self.
 socket = sockets
@@ -1035,215 +1035,25 @@
     
-killer = self.killer.connect_receiver()%0A%0A        while True:%0A            pollin, _, _ = zmq.select(%5Bsocket.socket, killer%5D, %5B%5D, %5B%5D)%0A            if pollin:%0A                if socket.socket in pollin:%0A        
+def socket_cb():%0A
     
@@ -1073,16 +1073,21 @@
 , msg = 
+self.
 socket.r
@@ -1108,29 +1108,13 @@
     
-      
-  
 try:%0A
-        
     
@@ -1189,24 +1189,16 @@
         
-        
 except A
@@ -1228,24 +1228,16 @@
         
- 
-       
 # TODO h
@@ -1259,24 +1259,16 @@
 ommands%0A
-        
         
@@ -1323,32 +1323,42 @@
 mmand')%0A        
+    else:%0A
                 
@@ -1361,197 +1361,472 @@
     
-pass%0A                    else:%0A                        handler(socket, ident, msg)%0A                if killer in pollin:%0A                    break%0A%0A        socket.close()%0A        killer.close
+handler(self.socket, ident, msg)%0A        self.register(socket, socket_cb)%0A%0A        killer = self.killer.connect_receiver()%0A        def killer_cb():%0A            killer.recv()%0A            for socket in list(self.sockets.keys()):%0A                socket.close()%0A                self.unregister(socket)%0A        self.register(killer, killer_cb)%0A%0A        while len(self._poller%5B1%5D) %3E 0:%0A            for socket, _ in self.poller.poll():%0A                self.sockets%5Bsocket%5D
 ()%0A

use a dict of callbacks instead of listing all relevant sockets in a series of ifs