From dd4b8b0eb2ead66e449c93ae9c5222288774ad2a Mon Sep 17 00:00:00 2001 From: Artur Borecki Date: Thu, 28 Aug 2025 13:10:05 +0200 Subject: [PATCH] feat(message.py): add `Message` class --- src/judas_protocol/message.py | 108 ++++++++++++++++++++++++++++++++++ 1 file changed, 108 insertions(+) create mode 100644 src/judas_protocol/message.py diff --git a/src/judas_protocol/message.py b/src/judas_protocol/message.py new file mode 100644 index 0000000..c702958 --- /dev/null +++ b/src/judas_protocol/message.py @@ -0,0 +1,108 @@ +# -*- coding: utf-8 -*- +from __future__ import annotations + +from enum import Enum +import json +from typing import Any + +from judas_protocol.types import Category, ControlAction + + +class Message: + """A message in the judas protocol.""" + + def __init__( + self, + category: Category, + action: Enum, + payload: dict[str, Any] | None = None, + ) -> None: + """Initialize a Message. + + Args: + category (Category): The category of the message. + action (Enum): The action of the message. + payload (dict[str, Any] | None): The payload of the message. + """ + self.category: Category = category + self.action: Enum = action + self.payload: dict[str, Any] = payload or {} + + def to_dict(self) -> dict[str, Any]: + """Convert the message to a dictionary. + + Returns: + dict[str, Any]: The message as a dictionary. + """ + return { + "category": self.category, + "action": self.action, + "payload": self.payload, + } + + def to_json(self) -> str: + """Convert the message to a JSON string. + + Returns: + str: The message as a JSON string. + """ + return json.dumps(self.to_dict()) + + @classmethod + def from_dict(cls, data: dict[str, Any]) -> Message: + """Create a Message from a dictionary. + + Args: + data (dict[str, Any]): The dictionary to create the message from. + + Returns: + Message: The created message. + """ + category = Category(data["category"]) + action_str = data["action"] + + match category: + case Category.CONTROL: + action = ControlAction(action_str) + case _: + raise ValueError(f"Unknown category: {category}") + + payload = data.get("payload", {}) + return cls(category=category, action=action, payload=payload) + + @classmethod + def from_json(cls, data: str) -> Message: + """Create a Message from a JSON string. + + Args: + data (str): The JSON string to create the message from. + Returns: + Message: The created message. + """ + return cls.from_dict(json.loads(data)) + + @classmethod + def ack(cls) -> Message: + """Create an ACK message. + + Returns: + Message: The created ACK message. + """ + return cls( + category=Category.CONTROL, action=ControlAction.ACK, payload={} + ) + + @classmethod + def hello(cls, mac: str) -> Message: + """Create a HELLO message. + + Args: + mac (str): The MAC address to include in the HELLO message. + Returns: + Message: The created HELLO message. + """ + return cls( + category=Category.CONTROL, + action=ControlAction.HELLO, + payload={"mac": mac}, + )