baseimage-tutorial-framework/src/components/component_base.py

import json
import zmq
from zmq.eventloop import ioloop
from zmq.eventloop.zmqstream import ZMQStream

from config import PUBLISHER_PORT, RECEIVER_PORT

ioloop.install()


class ComponentBase:
    def __init__(self, anchor, event_handler, zmq_context=None):
        self.anchor = anchor
        self.event_handler = event_handler
        self.zmq_context = zmq_context or zmq.Context.instance()
        self.zmq_sub_socket = self.zmq_context.socket(zmq.SUB)
        self.subscriptions = {self.anchor}
        self.zmq_sub_socket.setsockopt_string(zmq.SUBSCRIBE, self.anchor)
        self.zmq_sub_socket.connect('tcp://localhost:{}'.format(PUBLISHER_PORT))
        self.zmq_sub_stream = ZMQStream(self.zmq_sub_socket)
        self.zmq_push_socket = self.zmq_context.socket(zmq.PUSH)
        self.zmq_push_socket.connect('tcp://localhost:{}'.format(RECEIVER_PORT))

    def message_other(self, anchor, data):
        encoded_anchor = anchor.encode('utf-8')
        message = {
            'anchor': anchor,
            'data': data
        }
        encoded_message = json.dumps(message).encode('utf-8')
        self.zmq_push_socket.send_multipart([encoded_anchor, encoded_message])

    def subscribe(self, anchor):
        if anchor not in self.subscriptions:
            self.subscriptions.add(anchor)
            self.zmq_sub_socket.setsockopt_string(zmq.SUBSCRIBE, anchor)

    def unsubscribe(self, anchor):
        try:
            self.subscriptions.remove(anchor)
            self.zmq_sub_socket.setsockopt_string(zmq.UNSUBSCRIBE, anchor)
        except KeyError:
            pass

    def unsubscribe_all(self):
        for sub in self.subscriptions:
            self.zmq_sub_socket.setsockopt_string(zmq.UNSUBSCRIBE, sub)
        self.subscriptions.clear()