# zmq_reqresp.py import asyncio import zmq import zmq.asyncio from .background import Background from inspect import iscoroutinefunction class ZmqRequester(object): def __init__(self, url, async_mode=False, timeout=0): super().__init__() self.async_mode = async_mode self.url = url self.timeout = timeout self._connect() def __del__(self): self._close() def _connect(self): if self.async_mode: self.ctx = zmq.asyncio.Context() else: self.ctx = zmq.Context() # Socket to talk to server self.sock = self.ctx.socket(zmq.REQ) self.sock.connect(self.url) if self.timeout > 0: self.sock.setsockopt(zmq.LINGER, 0) self.poller = zmq.Poller() self.poller.register(self.sock, zmq.POLLIN) def _close(self): self.sock.close() self.ctx.term() def send(self, msg): """ send s string to responser, and return a string """ if self.async_mode: raise Exception('ZMQ_Requester: in async mode, use asend instead') b = msg.encode('utf-8') r = self.send_b(b) if r is not None: return r.decode('utf-8') return None def send_b(self, b): """ send a bytes and return a bytes """ if self.async_mode: raise Exception('ZMQ_Requester: in async mode, use asend_b instead') self.sock.send(b) if self.timeout > 0: if self.poller.poll(self.timeout * 1000): return self.sock.recv() else: self._close() self._connect() return None else: return self.sock.recv() async def asend_b(self, b): if not self.async_mode: raise Exception('ZMQ_Requester: not in async mode, use send_b instead') await self.sock.send_multipart([b]) if self.timeout > 0: if self.poller.poll(self.timeout * 1000): r = await self.sock.recv_multipart() return r else: self._close() self._connect() return None r = await self.sock.recv_multipart() return r[0] async def asend(self, msg): if not self.async_mode: raise Exception('ZMQ_Requester: not in async mode, use send instead') b = msg.encode('utf-8') r = await self.asend_b(b) if r is None: return None return r.decode('utf-8') class ZmqReplier(object): def __init__(self, url, handler, async_mode=False): self.async_mode = async_mode self.url = url if not self.async_mode and iscoroutinefunction(handler): raise('not in async mode, handler can not be a coroutine') self.handler = handler if self.async_mode: self.ctx = zmq.asyncio.Context() else: self.ctx = zmq.Context() self.sock = self.ctx.socket(zmq.REP) self.sock.bind(self.url) self.keep_running = True async def async_run(self): while self.keep_running: bs = await self.sock.recv_multipart() b = b[0] if iscoroutinefunction(self.handler): rb = await self.handler(b) else: rb =self.self.handler(b) if isinstance(rb, str): rb = rb.encode('utf-8') await self.sock.send_multipart([rb]) def run(self): self.background = Background(self._run) self.background.daemon = True self.background.start() def _run(self): while self.keep_running: b = self.sock.recv() rb = self.handler(b) if isinstance(rb, str): rb = rb.encode('utf-8') self.sock.send(rb) def stop(self): self.keep_running = False self.join()