baseimage-tutorial-framework/lib/tfw/event_handler_base.py

from abc import ABC, abstractmethod

from tfw.networking.serialization import deserialize_all
from tfw.networking.event_handlers.server_connector import ServerConnector


class EventHandlerBase(ABC):
    def __init__(self, key):
        self.server_connector = ServerConnector()
        self.key = key
        self.subscriptions = set()
        self.subscribe(self.key)
        self.subscribe('reset')
        self.server_connector.register_callback(self.event_handler_callback)

    def event_handler_callback(self, msg_parts):
        key, message = deserialize_all(*msg_parts)
        response = self.dispatch_handling(key, message)
        if response is None:
            return
        self.server_connector.send(key, response)

    @abstractmethod
    def dispatch_handling(self, key, message):
        raise NotImplementedError

    def _dispatch_handling(self, key, message):
        # pylint: disable=no-else-return
        if key != 'reset':
            return self.handle_event(key, message)
        else:
            return self.handle_reset(message)

    @abstractmethod
    def handle_event(self, key, message):
        raise NotImplementedError

    def handle_reset(self, message):
        # pylint: disable=unused-argument,no-self-use
        return None

    def cleanup(self):
        pass

    def message_other(self, key, data):
        message = {
            'key': key,
            'data': data
        }
        self.server_connector.send(key, message)

    def subscribe(self, key):
        if key not in self.subscriptions:
            self.subscriptions.add(key)
            self.server_connector.subscribe(key)

    def unsubscribe(self, key):
        try:
            self.subscriptions.remove(key)
            self.server_connector.unsubscribe(key)
        except KeyError:
            pass

    def unsubscribe_all(self):
        for sub in self.subscriptions:
            self.server_connector.unsubscribe(key=sub)
        self.subscriptions.clear()


class TriggerlessEventHandler(EventHandlerBase, ABC):
    # pylint: disable=abstract-method
    def dispatch_handling(self, key, message):
        return self._dispatch_handling(key, message)


class TriggeredEventHandler(EventHandlerBase, ABC):
    # pylint: disable=abstract-method
    def __init__(self, key, trigger):
        super().__init__(key)
        self.trigger = trigger

    def dispatch_handling(self, key, message):
        if message.get('trigger') == self.trigger:
            return self._dispatch_handling(key, message)
        return None