baseimage-tutorial-framework/lib/tfw/builtins/frontend_event_handler.py

from abc import ABC, abstractmethod
from contextlib import suppress

from tfw.networking import Scope

from .message_sender import MessageSender
from .event_handler import EventHandler


class FrontendEventHandler(EventHandler):
    def __init__(self):
        frontend_keys = ('message', 'queueMessages', 'dashboard', 'console')
        self._frontend_message_storage = FrontendMessageStorage(frontend_keys)
        super().__init__((*frontend_keys, 'recover'), scope=Scope.WEBSOCKET)

    def handle_event(self, message):
        self._frontend_message_storage.save_message(message)
        if message['key'] == 'recover':
            self.recover_frontend()
        self.send_message(message)

    def recover_frontend(self):
        for message in self._frontend_message_storage.messages:
            self.send_message(message)


class MessageStorage(ABC):
    def __init__(self):
        self._messages = []

    def save_message(self, message):
        with suppress(KeyError, AttributeError):
            if self._filter_message(message):
                self._messages.extend(self._transform_message(message))

    @abstractmethod
    def _filter_message(self, message):
        raise NotImplementedError

    def _transform_message(self, message):   # pylint: disable=no-self-use
        yield message

    def clear(self):
        self._messages.clear()

    @property
    def messages(self):
        yield from self._messages


class FrontendMessageStorage(MessageStorage):
    def __init__(self, keys):
        self._keys = keys
        super().__init__()

    def _filter_message(self, message):
        key = message['key']
        return key in self._keys

    def _transform_message(self, message):
        if message['key'] == 'queueMessages':
            yield from MessageSender.generate_messages_from_queue(message)
        else:
            yield message