message mux tool for python

multiplex.py

from dataclasses import dataclass, field, ClassVar
from typing import Callable, Union, Optional, Any
from enum import Enum

Pattern = Union[Enum, str, int]
Message = Any

MessageHandler = Callable[[Message], None]
MessageDecorator = Callable[[MessageHandler], MessageHandler]

@dataclass
class NotFoundError(Exception):
    message: ClassVar[str] = "not found"
    http_code: ClassVar[int] = 404

def link(handler: MessageHandler, *decorators: MessageDecorator) -> MessageHandler:
    for decorator in reversed(decorators):
        handler = decorator(handler)
    return handler

@dataclass
class MessageMux:
    get_pattern: Callable[[Message], Pattern]
    not_found_handler: Optional[MessageHandler] = None
    handlers: dict[Pattern, MessageHandler] = field(default_factory=dict)
    decorators: list[MessageDecorator] = field(default_factory=list)

    def handle_message(self, dto: Message):
        pattern = self.get_pattern(dto)

        handler = self.handlers.get(pattern)
        if not handler:
            if not self.not_found_handler:
                raise NotFoundError(f"pattern={pattern} does not exist in mux")
            self.not_found_handler(dto)

        link(handler, *self.decorators)(dto)

    def register_handler(self, pattern: Pattern, handler: MessageHandler) -> 'MessageMux':
        if pattern in self.handlers:
            raise ValueError(f"mux has duplicate pattern={pattern}")
        self.handlers[pattern] = handler
        return self

    def remove_handler(self, pattern: Pattern) -> None:
        if pattern in self.handlers:
            del self.handlers[pattern]

    def add_middleware(self, *middlewares: MessageDecorator) -> 'MessageMux':
        self.decorators.extend(middlewares)
        return self

    def set_not_found_handler(self, handler: MessageHandler) -> 'MessageMux':
        self.not_found_handler = handler
        return self