Source

pyzmq-article / example_app / pongproc.py

Full commit
from zmq.utils import jsonapi as json
import zmq

import zmqproc


host = '127.0.0.1'
port = 5678


def ping():
    """Sends ping requests and waits for replies."""
    context = zmq.Context()
    sock = context.socket(zmq.REQ)
    sock.connect('tcp://%s:%s' % (host, port))

    for i in range(5):
        sock.send_json(['ping', i])
        rep = sock.recv_json()
        print('Ping got reply:', rep)

    sock.send_json(['plzdiekthxbye', None])


class PongProc(zmqproc.ZmqProcess):
    """
    Main processes for the Ponger. It handles ping requests and sends back
    a pong.

    """
    def __init__(self, bind_addr):
        super().__init__()

        self.bind_addr = bind_addr
        self.rep_stream = None

        # Make sure this is pickle-able (e.g., not using threads)
        # or it won't work on Windows. If it's not pickle-able, instantiate
        # it in setup().
        self.ping_handler = PingHandler()

    def setup(self):
        """Sets up PyZMQ and creates all streams."""
        super().setup()

        self.rep_stream, _ = self.stream(zmq.REP, self.bind_addr, bind=True,
                callback=self.handle_rep_stream)

    def run(self):
        """Sets up everything and starts the event loop."""
        self.setup()
        self.loop.start()

    def stop(self):
        """Stops the event loop."""
        self.loop.stop()

    def handle_rep_stream(self, msg):
        """
        Handles messages from a Pinger:

        *ping*
            Send back a pong.

        *plzdiekthxbye*
            Stop the ioloop and exit.

        """
        msg_type, data = json.loads(msg[0])

        if msg_type == 'ping':
            rep = self.ping_handler.make_pong(data)
            self.rep_stream.send_json(rep)

        elif msg_type == 'plzdiekthxbye':
            self.stop()

        else:
            raise RuntimeError('Received unkown message type: %s' % msg_type)


class PingHandler(object):

    def make_pong(self, num_pings):
        """Creates and returns a pong message."""
        print('Pong got request number %s' % num_pings)

        return ['pong', num_pings]


if __name__ == '__main__':
    pong_proc = PongProc(bind_addr=(host, port))
    pong_proc.start()

    ping()

    pong_proc.join()