From 3ab522baf885525ddeb8be94bde0026b8e04f769 Mon Sep 17 00:00:00 2001
From: Felix Stupp <felix.stupp@banananet.work>
Date: Sat, 24 Feb 2024 23:21:35 +0000
Subject: [PATCH] comm/transport/datagrams: add msg_type map for deserialize

---
 enjo_lib/comm/transport/datagrams.py | 13 +++++++++++++
 1 file changed, 13 insertions(+)

diff --git a/enjo_lib/comm/transport/datagrams.py b/enjo_lib/comm/transport/datagrams.py
index eba744e..1edf8bc 100644
--- a/enjo_lib/comm/transport/datagrams.py
+++ b/enjo_lib/comm/transport/datagrams.py
@@ -1,10 +1,12 @@
 from __future__ import annotations
 
 from collections.abc import (
+    Mapping,
     Sequence,
 )
 from typing import (
     Literal,
+    Type,
 )
 
 from attrs import (
@@ -159,3 +161,14 @@ class RestartDatagram(GeneralTransportDatagram):
 class TransmitDatagram(GeneralTransportDatagram):
     message_type: Literal["transmit"] = "transmit"
     data: JSON
+
+
+MESSAGE_TYPE_MAP: Mapping[str, Type[GeneralTransportDatagram]] = {
+    "connect": ConnectDatagram,
+    "disconnect": DisconnectDatagram,
+    "duplicate_sender": DuplicateSenderDatagram,
+    "replace_sender": ReplaceSenderDatagram,
+    "failure": FailureDatagram,
+    "restart": RestartDatagram,
+    "transmit": TransmitDatagram,
+}