def zmq_import_poll(sock: zmq.Socket, timeout: float): """ Author: Alexander Heilmeier Description: Handles incoming ZMQ messages and allows polling. Inputs: sock: ZMQ socket (see below how to create it) timeout: [s] timeout for polling How to create a ZMQ socket to import data? import zmq zmq_context = zmq.Context() sock = zmq_context.socket(zmq.PUB) sock.connect("tcp://%s:%s" % (ip, port)) sock.setsockopt_string(zmq.SUBSCRIBE, zmq_topic]) """ # ------------------------------------------------------------------------------------------------------------------ # PREPARATION ---------------------------------------------------------------------------------------------------- # ------------------------------------------------------------------------------------------------------------------ # timeout conversion s -> ms timeout *= 1000.0 # set standard return data = None # ------------------------------------------------------------------------------------------------------------------ # CHECK IF BUFFER ALREADY STORES VALUES ---------------------------------------------------------------------------- # ------------------------------------------------------------------------------------------------------------------ # non blocking socket to empty buffer try: while True: sock.recv_string(flags=zmq.NOBLOCK) data = sock.recv_pyobj(flags=zmq.NOBLOCK) except zmq.Again: pass # ------------------------------------------------------------------------------------------------------------------ # WAIT FOR A SHORT TIME IF NO DATA WAS STORED IN THE BUFFER -------------------------------------------------------- # ------------------------------------------------------------------------------------------------------------------ if data is None: # use poll to be able to set a timeout -> will return 1 if an event occurs, otherwise 0 events = sock.poll(timeout=timeout) if events: # non blocking socket to empty buffer try: while True: sock.recv_string(flags=zmq.NOBLOCK) data = sock.recv_pyobj(flags=zmq.NOBLOCK) except zmq.Again: pass return data
def recv_json(socket: zmq.Socket, **kw) -> Dict: """ Receive a dictionary. """ s_data = socket.recv_string(**kw) return json.loads(s_data)
def zmq_import(sock: zmq.Socket, blocking: bool = False, datatype: str = "pyobj", use_buffer: bool = False): """ Author: Alexander Heilmeier & Tim Stahl Description: Handles incoming ZMQ messages. Inputs: sock: ZMQ socket (see below how to create it) blocking: set if socket should be blocking (i.e. wait until new message is received) or not datatype: string that indicates if it should be received as Python object (pyobj), json (json) or string (str) use_buffer: flag to indicate if the buffer should be returned completely (True) or only the latest message (False) Hint: Conversion of json objects to their original data type is handled by PyZMQ and therefore must not be done by hand. How to create a ZMQ socket to import data? import zmq zmq_context = zmq.Context() sock = zmq_context.socket(zmq.PUB) sock.connect("tcp://%s:%s" % (ip, port)) sock.setsockopt_string(zmq.SUBSCRIBE, zmq_topic]) """ # ------------------------------------------------------------------------------------------------------------------ # Specify called receive handle ------------------------------------------------------------------------------------ # ------------------------------------------------------------------------------------------------------------------ if datatype == "pyobj": sock_recv_fct = lambda **kwargs: sock.recv_pyobj(**kwargs) elif datatype == "json": sock_recv_fct = lambda **kwargs: sock.recv_json(**kwargs) elif datatype == "str": sock_recv_fct = lambda **kwargs: sock.recv_string(**kwargs) else: raise RuntimeError("Specified datatype is not supported!") # ------------------------------------------------------------------------------------------------------------------ # FUNCTION BODY ---------------------------------------------------------------------------------------------------- # ------------------------------------------------------------------------------------------------------------------ data = None buffer = [] # non blocking socket to empty buffer try: while True: sock.recv_string(flags=zmq.NOBLOCK) data = sock_recv_fct(flags=zmq.NOBLOCK) if use_buffer: buffer.append(data) except zmq.Again: pass # if no data was received above and blocking is True, wait for new data if blocking and not use_buffer and data is None: sock.recv_string() data = sock_recv_fct() elif blocking and use_buffer and not buffer: sock.recv_string() data = sock_recv_fct() buffer.append(data) if use_buffer and buffer: return buffer elif use_buffer: return None else: return data
def recv_string(socket: zmq.Socket, **kw) -> str: """ Receive a string. """ return socket.recv_string(**kw)