Python zmq.Socket() Examples

def _set_handler(self, socket, handler, update=False):
        """
        Set the socket handler(s).

        Parameters
        ----------
        socket : zmq.Socket
            Socket to set its handler(s).
        handler : function(s)
            Handler(s) for the socket. This can be a list or a dictionary too.
        """
        if update:
            try:
                self._handler[socket].update(self._curated_handlers(handler))
            except KeyError:
                self._handler[socket] = self._curated_handlers(handler)
        else:
            self._handler[socket] = self._curated_handlers(handler) 

def test_shadow_pyczmq(self):
        try:
            from pyczmq import zctx, zsocket
        except Exception:
            raise SkipTest("Requires pyczmq")
        
        ctx = zctx.new()
        ca = zsocket.new(ctx, zmq.PUSH)
        cb = zsocket.new(ctx, zmq.PULL)
        a = zmq.Socket.shadow(ca)
        b = zmq.Socket.shadow(cb)
        a.bind("inproc://a")
        b.connect("inproc://a")
        a.send(b'hi')
        rcvd = self.recv(b)
        self.assertEqual(rcvd, b'hi') 

def _recv_ack(sub_channel: Socket, topic: str, expected_content: bool = True):
    """ Blocks until an acknowledge-message for the specified topic with the expected content is received via the
        specified subscriber channel. 
    
    Args:
        sub_channel (Socket): subscriber socket
        topic (str): topic to listen for ACK's
        expected_content (bool): are we expecting `True` (ACK) or `False` (NACK)
    """
    ack_topic = topic if topic.startswith("ACK/") else f"ACK/{topic}"
    while True:
        msg = sub_channel.recv_multipart(copy=True)
        recv_topic = msg[0].decode("ascii")
        content = pickle.loads(msg[1])[1]  # pickle.loads(msg) -> tuple(timestamp, content) -> return content
        if recv_topic == ack_topic:
            if content == expected_content:
                return 

def test_skt_reinit():
    result = {'foo': None, 'bar': None}

    @socket(zmq.PUB)
    def f(key, skt):
        assert isinstance(skt, zmq.Socket), skt

        result[key] = skt

    foo_t = threading.Thread(target=f, args=('foo',))
    bar_t = threading.Thread(target=f, args=('bar',))

    foo_t.start()
    bar_t.start()

    foo_t.join()
    bar_t.join()

    assert result['foo'] is not None, result
    assert result['bar'] is not None, result
    assert result['foo'] is not result['bar'], result 

def test_shadow(self):
        p = self.socket(zmq.PUSH)
        p.bind("tcp://127.0.0.1:5555")
        p2 = zmq.Socket.shadow(p.underlying)
        self.assertEqual(p.underlying, p2.underlying)
        s = self.socket(zmq.PULL)
        s2 = zmq.Socket.shadow(s.underlying)
        self.assertNotEqual(s.underlying, p.underlying)
        self.assertEqual(s.underlying, s2.underlying)
        s2.connect("tcp://127.0.0.1:5555")
        sent = b'hi'
        p2.send(sent)
        rcvd = self.recv(s2)
        self.assertEqual(rcvd, sent) 

def test_shadow_pyczmq(self):
        try:
            from pyczmq import zctx, zsocket
        except Exception:
            raise SkipTest("Requires pyczmq")
        
        ctx = zctx.new()
        ca = zsocket.new(ctx, zmq.PUSH)
        cb = zsocket.new(ctx, zmq.PULL)
        a = zmq.Socket.shadow(ca)
        b = zmq.Socket.shadow(cb)
        a.bind("inproc://a")
        b.connect("inproc://a")
        a.send(b'hi')
        rcvd = self.recv(b)
        self.assertEqual(rcvd, b'hi') 

def __init__(self, context=None, socket_type=-1, io_loop=None, **kwargs):
        if isinstance(context, _zmq.Socket):
            context, from_socket = (None, context)
        else:
            from_socket = kwargs.pop('_from_socket', None)
        if from_socket is not None:
            super(_AsyncSocket, self).__init__(shadow=from_socket.underlying)
            self._shadow_sock = from_socket
        else:
            super(_AsyncSocket, self).__init__(context, socket_type, **kwargs)
            self._shadow_sock = _zmq.Socket.shadow(self.underlying)
        self.io_loop = io_loop or self._default_loop()
        self._recv_futures = deque()
        self._send_futures = deque()
        self._state = 0
        self._fd = self._shadow_sock.FD
        self._init_io_state() 

def poll(self, timeout=None):
        """Poll the registered 0MQ or native fds for I/O.

        Parameters
        ----------
        timeout : float, int
            The timeout in milliseconds. If None, no `timeout` (infinite). This
            is in milliseconds to be compatible with ``select.poll()``.

        Returns
        -------
        events : list of tuples
            The list of events that are ready to be processed.
            This is a list of tuples of the form ``(socket, event)``, where the 0MQ Socket
            or integer fd is the first element, and the poll event mask (POLLIN, POLLOUT) is the second.
            It is common to call ``events = dict(poller.poll())``,
            which turns the list of tuples into a mapping of ``socket : event``.
        """
        if timeout is None or timeout < 0:
            timeout = -1
        elif isinstance(timeout, float):
            timeout = int(timeout)
        return zmq_poll(self.sockets, timeout=timeout) 

def _create_socket(self, socket_type, linger_value):
        """
        Create a socket of the given type, bind it to a random port and
        register it to the poller

        :param int socket_type: type of the socket to open
        :param int linger_value: -1 mean wait for receive all msg and block
                                 closing 0 mean hardkill the socket even if msg
                                 are still here.
        :return (zmq.Socket, int): the initialized socket and the port where the
                                   socket is bound
        """
        socket = SafeContext.get_context().socket(socket_type)
        socket.setsockopt(zmq.LINGER, linger_value)
        socket.set_hwm(0)
        port_number = socket.bind_to_random_port(LOCAL_ADDR)
        self.poller.register(socket, zmq.POLLIN)
        self.logger.debug("bind to " + LOCAL_ADDR + ':' + str(port_number))
        return (socket, port_number) 

def test_multi_skts():
    @socket(zmq.PUB)
    @socket(zmq.SUB)
    @socket(zmq.PUSH)
    def test(pub, sub, push):
        assert isinstance(pub, zmq.Socket), pub
        assert isinstance(sub, zmq.Socket), sub
        assert isinstance(push, zmq.Socket), push

        assert pub.context is zmq.Context.instance()
        assert sub.context is zmq.Context.instance()
        assert push.context is zmq.Context.instance()

        assert pub.type == zmq.PUB
        assert sub.type == zmq.SUB
        assert push.type == zmq.PUSH
    test() 

def test_skt_reinit():
    result = {'foo': None, 'bar': None}

    @socket(zmq.PUB)
    def f(key, skt):
        assert isinstance(skt, zmq.Socket), skt

        result[key] = skt

    foo_t = threading.Thread(target=f, args=('foo',))
    bar_t = threading.Thread(target=f, args=('bar',))

    foo_t.start()
    bar_t.start()

    foo_t.join()
    bar_t.join()

    assert result['foo'] is not None, result
    assert result['bar'] is not None, result
    assert result['foo'] is not result['bar'], result 

def test_multi_skts_with_name():
    @socket('foo', zmq.PUSH)
    @socket('bar', zmq.SUB)
    @socket('baz', zmq.PUB)
    def test(foo, bar, baz):
        assert isinstance(foo, zmq.Socket), foo
        assert isinstance(bar, zmq.Socket), bar
        assert isinstance(baz, zmq.Socket), baz

        assert foo.context is zmq.Context.instance()
        assert bar.context is zmq.Context.instance()
        assert baz.context is zmq.Context.instance()

        assert foo.type == zmq.PUSH
        assert bar.type == zmq.SUB
        assert baz.type == zmq.PUB
    test() 

def test_multi_skts_with_name():
    @socket('foo', zmq.PUSH)
    @socket('bar', zmq.SUB)
    @socket('baz', zmq.PUB)
    def test(foo, bar, baz):
        assert isinstance(foo, zmq.Socket), foo
        assert isinstance(bar, zmq.Socket), bar
        assert isinstance(baz, zmq.Socket), baz

        assert foo.context is zmq.Context.instance()
        assert bar.context is zmq.Context.instance()
        assert baz.context is zmq.Context.instance()

        assert foo.type == zmq.PUSH
        assert bar.type == zmq.SUB
        assert baz.type == zmq.PUB
    test() 

def test_skt_multi_thread():
    @socket(zmq.PUB)
    @socket(zmq.SUB)
    @socket(zmq.PUSH)
    def f(pub, sub, push):
        assert isinstance(pub, zmq.Socket), pub
        assert isinstance(sub, zmq.Socket), sub
        assert isinstance(push, zmq.Socket), push

        assert pub.context is zmq.Context.instance()
        assert sub.context is zmq.Context.instance()
        assert push.context is zmq.Context.instance()

        assert pub.type == zmq.PUB
        assert sub.type == zmq.SUB
        assert push.type == zmq.PUSH

        assert len(set(map(id, [pub, sub, push]))) == 3

    threads = [threading.Thread(target=f) for i in range(8)]
    [t.start() for t in threads]
    [t.join() for t in threads] 

def test_skt_multi_thread():
    @socket(zmq.PUB)
    @socket(zmq.SUB)
    @socket(zmq.PUSH)
    def f(pub, sub, push):
        assert isinstance(pub, zmq.Socket), pub
        assert isinstance(sub, zmq.Socket), sub
        assert isinstance(push, zmq.Socket), push

        assert pub.context is zmq.Context.instance()
        assert sub.context is zmq.Context.instance()
        assert push.context is zmq.Context.instance()

        assert pub.type == zmq.PUB
        assert sub.type == zmq.SUB
        assert push.type == zmq.PUSH

        assert len(set(map(id, [pub, sub, push]))) == 3

    threads = [threading.Thread(target=f) for i in range(8)]
    [t.start() for t in threads]
    [t.join() for t in threads] 

def multi_skts_method_other_args(self):
        @socket(zmq.PUB)
        @socket(zmq.SUB)
        def f(foo, pub, sub, bar=None):
            assert isinstance(pub, zmq.Socket), pub
            assert isinstance(sub, zmq.Socket), sub

            assert foo == 'mock'
            assert bar == 'fake'

            assert pub.context is zmq.Context.instance()
            assert sub.context is zmq.Context.instance()

            assert pub.type is zmq.PUB
            assert sub.type is zmq.SUB

        f('mock', bar='fake') 

def _process_pull_event(self, socket, addr, data):
        """
        Process a PULL socket's event.

        Parameters
        ----------
        socket : zmq.Socket
            Socket that generated the event.
        addr : AgentAddress
            AgentAddress associated with the socket that generated the event.
        data : bytes
            Data received on the socket.
        """
        message = deserialize_message(message=data, serializer=addr.serializer)
        handler = self._handler[socket]
        if not isinstance(handler, (list, dict, tuple)):
            handler = [handler]
        for h in handler:
            h(self, message) 

def multi_skts_method_other_args(self):
        @socket(zmq.PUB)
        @socket(zmq.SUB)
        def f(foo, pub, sub, bar=None):
            assert isinstance(pub, zmq.Socket), pub
            assert isinstance(sub, zmq.Socket), sub

            assert foo == 'mock'
            assert bar == 'fake'

            assert pub.context is zmq.Context.instance()
            assert sub.context is zmq.Context.instance()

            assert pub.type is zmq.PUB
            assert sub.type is zmq.SUB

        f('mock', bar='fake') 

def _process_rep_event(self, socket, addr, data):
        """
        Process a REP socket's event.

        Parameters
        ----------
        socket : zmq.Socket
            Socket that generated the event.
        addr : AgentAddress
            AgentAddress associated with the socket that generated the event.
        data : bytes
            Data received on the socket.
        """
        message = deserialize_message(message=data, serializer=addr.serializer)
        handler = self._handler[socket]
        if inspect.isgeneratorfunction(handler):
            generator = handler(self, message)
            socket.send(serialize_message(next(generator), addr.serializer))
            execute_code_after_yield(generator)
        else:
            reply = handler(self, message)
            socket.send(serialize_message(reply, addr.serializer)) 

def poll(self, timeout=None):
        """Poll the registered 0MQ or native fds for I/O.

        Parameters
        ----------
        timeout : float, int
            The timeout in milliseconds. If None, no `timeout` (infinite). This
            is in milliseconds to be compatible with ``select.poll()``. The
            underlying zmq_poll uses microseconds and we convert to that in
            this function.
        
        Returns
        -------
        events : list of tuples
            The list of events that are ready to be processed.
            This is a list of tuples of the form ``(socket, event)``, where the 0MQ Socket
            or integer fd is the first element, and the poll event mask (POLLIN, POLLOUT) is the second.
            It is common to call ``events = dict(poller.poll())``,
            which turns the list of tuples into a mapping of ``socket : event``.
        """
        if timeout is None or timeout < 0:
            timeout = -1
        elif isinstance(timeout, float):
            timeout = int(timeout)
        return zmq_poll(self.sockets, timeout=timeout) 

def _process_single_event(self, socket):
        """
        Process a socket's event.

        Parameters
        ----------
        socket : zmq.Socket
            Socket that generated the event.
        """
        data = socket.recv()
        address = self._address[socket]
        if address.kind == 'SUB':
            self._process_sub_event(socket, address, data)
        elif address.kind == 'PULL':
            self._process_pull_event(socket, address, data)
        elif address.kind == 'REP':
            self._process_rep_event(socket, address, data)
        else:
            self._process_single_event_complex(address, socket, data) 

def test_multi_skts():
    @socket(zmq.PUB)
    @socket(zmq.SUB)
    @socket(zmq.PUSH)
    def test(pub, sub, push):
        assert isinstance(pub, zmq.Socket), pub
        assert isinstance(sub, zmq.Socket), sub
        assert isinstance(push, zmq.Socket), push

        assert pub.context is zmq.Context.instance()
        assert sub.context is zmq.Context.instance()
        assert push.context is zmq.Context.instance()

        assert pub.type == zmq.PUB
        assert sub.type == zmq.SUB
        assert push.type == zmq.PUSH
    test() 

def test_ctx_skt_reinit():
    result = {'foo': {'ctx': None, 'skt': None},
              'bar': {'ctx': None, 'skt': None}}

    @context()
    @socket(zmq.PUB)
    def f(key, ctx, skt):
        assert isinstance(ctx, zmq.Context), ctx
        assert isinstance(skt, zmq.Socket), skt

        result[key]['ctx'] = ctx
        result[key]['skt'] = skt

    foo_t = threading.Thread(target=f, args=('foo',))
    bar_t = threading.Thread(target=f, args=('bar',))

    foo_t.start()
    bar_t.start()

    foo_t.join()
    bar_t.join()

    assert result['foo']['ctx'] is not None, result
    assert result['foo']['skt'] is not None, result
    assert result['bar']['ctx'] is not None, result
    assert result['bar']['skt'] is not None, result
    assert result['foo']['ctx'] is not result['bar']['ctx'], result
    assert result['foo']['skt'] is not result['bar']['skt'], result 

def _get_descriptors(self):
        """Returns three elements tuple with socket descriptors ready
        for gevent.select.select
        """
        rlist = []
        wlist = []
        xlist = []

        for socket, flags in self.sockets:
            if isinstance(socket, zmq.Socket):
                rlist.append(socket.getsockopt(zmq.FD))
                continue
            elif isinstance(socket, int):
                fd = socket
            elif hasattr(socket, 'fileno'):
                try:
                    fd = int(socket.fileno())
                except:
                    raise ValueError('fileno() must return an valid integer fd')
            else:
                raise TypeError('Socket must be a 0MQ socket, an integer fd '
                                'or have a fileno() method: %r' % socket)

            if flags & zmq.POLLIN:
                rlist.append(fd)
            if flags & zmq.POLLOUT:
                wlist.append(fd)
            if flags & zmq.POLLERR:
                xlist.append(fd)

        return (rlist, wlist, xlist) 

def _process_async_rep_event(self, socket, channel, data):
        """
        Process a ASYNC_REP socket's event.

        Parameters
        ----------
        socket : zmq.Socket
            Socket that generated the event.
        channel : AgentChannel
            AgentChannel associated with the socket that generated the event.
        data : bytes
            Data received on the socket.
        """
        message = deserialize_message(
            message=data, serializer=channel.serializer
        )
        address_uuid, request_uuid, data, address = message
        client_address = address.twin()
        if not self._registered(client_address):
            self.connect(address)
        handler = self._handler[socket]
        is_generator = inspect.isgeneratorfunction(handler)
        if is_generator:
            generator = handler(self, data)
            reply = next(generator)
        else:
            reply = handler(self, data)
        self.send(client_address, (address_uuid, request_uuid, reply))
        if is_generator:
            execute_code_after_yield(generator) 

def multi_skts_method(self, ctx, pub, sub, foo='bar'):
        assert isinstance(self, TestMethodDecorators), self
        assert isinstance(ctx, zmq.Context), ctx
        assert isinstance(pub, zmq.Socket), pub
        assert isinstance(sub, zmq.Socket), sub
        assert foo == 'bar'

        assert pub.context is ctx
        assert sub.context is ctx

        assert pub.type is zmq.PUB
        assert sub.type is zmq.SUB 

def test_skt_default_ctx():
    @socket(zmq.PUB)
    def test(skt):
        assert isinstance(skt, zmq.Socket), skt
        assert skt.context is zmq.Context.instance()
        assert skt.type == zmq.PUB
    test() 

def send_arrays(socket, arrays, stop=False):
  """Send NumPy arrays using the buffer interface and some metadata.

  Parameters
  ----------
  socket : :class:`zmq.Socket`
  The socket to send data over.
  arrays : list
  A list of :class:`numpy.ndarray` to transfer.
  stop : bool, optional
  Instead of sending a series of NumPy arrays, send a JSON object
  with a single `stop` key. The :func:`recv_arrays` will raise
  ``StopIteration`` when it receives this.

  Notes
  -----
  The protocol is very simple: A single JSON object describing the array
  format (using the same specification as ``.npy`` files) is sent first.
  Subsequently the arrays are sent as bytestreams (through NumPy's
  support of the buffering protocol).

  """
  if arrays:
    # The buffer protocol only works on contiguous arrays
    arrays = [numpy.ascontiguousarray(array) for array in arrays]
  if stop:
    headers = {'stop': True}
    socket.send_json(headers)
  else:
    headers = [header_data_from_array_1_0(array) for array in arrays]
    socket.send_json(headers, zmq.SNDMORE)
    for array in arrays[:-1]:
      socket.send(array, zmq.SNDMORE)
    socket.send(arrays[-1]) 

def recv_arrays(socket):
  """Receive a list of NumPy arrays.

  Parameters
  ----------
  socket : :class:`zmq.Socket`
  The socket to receive the arrays on.

  Returns
  -------
  list
  A list of :class:`numpy.ndarray` objects.

  Raises
  ------
  StopIteration
  If the first JSON object received contains the key `stop`,
  signifying that the server has finished a single epoch.

  """
  headers = socket.recv_json()
  if 'stop' in headers:
    raise StopIteration
  arrays = []
  for header in headers:
    data = socket.recv()
    buf = buffer_(data)
    array = numpy.frombuffer(buf, dtype=numpy.dtype(header['descr']))
    array.shape = header['shape']
    if header['fortran_order']:
      array.shape = header['shape'][::-1]
      array = array.transpose()
    arrays.append(array)
  return arrays 

def test_ctx_skt_name():
    @context('ctx')
    @socket('skt', zmq.PUB, context_name='ctx')
    def test(ctx, skt):
        assert isinstance(skt, zmq.Socket), skt
        assert isinstance(ctx, zmq.Context), ctx
        assert skt.type == zmq.PUB
    test()