fix(typed_state): fix support for data class in collections

Thibault-Pelletier · jourdain · commit 7e472f6d818e · 2025-09-19T11:36:35.000-06:00
- Fix encoding / decoding of dataclasses nested in lists or dicts
diff --git a/tests/test_typed_state.py b/tests/test_typed_state.py
@@ -437,3 +437,89 @@ def decode(self, _obj, _obj_type: type):
 
     with pytest.raises(TypeError):
         print(typed_state.data.my_enum)
+
+
+@dataclass
+class SimpleTypes:
+    my_int: int
+    my_enum: MyEnum
+    my_path: Path
+
+
+@dataclass
+class TypedComposite:
+    simple_types: SimpleTypes = field(default_factory=SimpleTypes)
+
+
+@dataclass
+class DataclassCollections:
+    nested_list: list[TypedComposite] = field(default_factory=list)
+    nested_dict: dict[str, TypedComposite] = field(default_factory=dict)
+
+
+def test_encode_decode_supports_collections_of_nested_dataclass(state):
+    typed_state = TypedState(state, DataclassCollections)
+    typed_state.data.nested_list = [
+        TypedComposite(
+            SimpleTypes(my_int=1, my_enum=MyEnum.A, my_path=Path("/path/to/1"))
+        ),
+        TypedComposite(
+            SimpleTypes(my_int=2, my_enum=MyEnum.B, my_path=Path("/path/to/2"))
+        ),
+    ]
+
+    typed_state.data.nested_dict = {
+        "3": TypedComposite(
+            SimpleTypes(my_int=3, my_enum=MyEnum.C, my_path=Path("/path/to/3"))
+        ),
+        "4": TypedComposite(
+            SimpleTypes(my_int=4, my_enum=MyEnum.A, my_path=Path("/path/to/4"))
+        ),
+    }
+
+    assert typed_state.data.nested_list[0] == TypedComposite(
+        SimpleTypes(my_int=1, my_enum=MyEnum.A, my_path=Path("/path/to/1"))
+    )
+    assert typed_state.data.nested_list[1] == TypedComposite(
+        SimpleTypes(my_int=2, my_enum=MyEnum.B, my_path=Path("/path/to/2"))
+    )
+    assert typed_state.data.nested_dict["3"] == TypedComposite(
+        SimpleTypes(my_int=3, my_enum=MyEnum.C, my_path=Path("/path/to/3"))
+    )
+    assert typed_state.data.nested_dict["4"] == TypedComposite(
+        SimpleTypes(my_int=4, my_enum=MyEnum.A, my_path=Path("/path/to/4"))
+    )
+
+    assert state[typed_state.name.nested_list] == [
+        {
+            "simple_types": {
+                "my_int": 1,
+                "my_enum": typed_state.encode(MyEnum.A),
+                "my_path": typed_state.encode("/path/to/1"),
+            }
+        },
+        {
+            "simple_types": {
+                "my_int": 2,
+                "my_enum": typed_state.encode(MyEnum.B),
+                "my_path": typed_state.encode("/path/to/2"),
+            }
+        },
+    ]
+
+    assert state[typed_state.name.nested_dict] == {
+        "3": {
+            "simple_types": {
+                "my_int": 3,
+                "my_enum": typed_state.encode(MyEnum.C),
+                "my_path": typed_state.encode("/path/to/3"),
+            }
+        },
+        "4": {
+            "simple_types": {
+                "my_int": 4,
+                "my_enum": typed_state.encode(MyEnum.A),
+                "my_path": typed_state.encode("/path/to/4"),
+            }
+        },
+    }
diff --git a/trame_server/utils/typed_state.py b/trame_server/utils/typed_state.py
@@ -115,6 +115,12 @@ def __init__(self, encoders: list[IStateEncoderDecoder] | None = None):
         self._encoders = encoders or [DefaultEncoderDecoder()]
 
     def encode(self, obj):
+        if is_dataclass(obj):
+            return {
+                field.name: self.encode(getattr(obj, field.name))
+                for field in fields(obj)
+            }
+
         if isinstance(obj, dict):
             return {self.encode(key): self.encode(value) for key, value in obj.items()}
 
@@ -156,6 +162,7 @@ def _try_decode(self, obj, obj_type: type):
 
     def _decode_strategies(self) -> list[Callable[[Any, type], Any]]:
         return [
+            self._decode_dataclass,
             self._decode_union,
             self._decode_dict,
             self._decode_iterable,
@@ -196,6 +203,17 @@ def _decode_union(self, obj, obj_type: type):
                 return val
         return self.failed_serialization()
 
+    def _decode_dataclass(self, obj, obj_type: type):
+        if not is_dataclass(obj_type):
+            return self.failed_serialization()
+
+        field_types = get_type_hints(obj_type)
+        decoded_dict = {
+            field.name: self._try_decode(obj.get(field.name), field_types[field.name])
+            for field in fields(obj_type)
+        }
+        return obj_type(**decoded_dict)
+
     @classmethod
     def _is_union_type(cls, obj_type: type):
         return get_origin(obj_type) is Union or isinstance(obj_type, UnionType)
