feat(v2): improve encoding and type checking

Author: jourdain

src/trame_dataclass/v2.py

@@ -5,6 +5,7 @@
 import weakref
 from collections.abc import Awaitable, Sequence
 from dataclasses import dataclass, field
+from enum import Enum, auto
 from typing import Any, Callable, Union, get_args, get_origin
 
 from loguru import logger
@@ -117,6 +118,17 @@ def _save_field(name, src, dst, encoder=None):
         dst[name] = value
 
 
+def _setup_class_fields(owner):
+    # set
+    for key in ["FIELD_NAMES", "DATACLASS_NAMES", "CLIENT_NAMES", "CLIENT_ONLY_NAMES"]:
+        if not hasattr(owner, key):
+            setattr(owner, key, set())
+    # dict
+    for key in ["ENCODERS", "TYPE_CHECKING"]:
+        if not hasattr(owner, key):
+            setattr(owner, key, {})
+
+
 # -----------------------------------------------------------------------------
 # Dataclass builder
 # -----------------------------------------------------------------------------
@@ -126,6 +138,12 @@ def __init__(self, encoder, decoder):
         self.decoder = decoder
 
 
+class TypeValidation(Enum):
+    STRICT = auto()
+    WARNING = auto()
+    SKIP = auto()
+
+
 class StateDataModel:
     def __init__(self, trame_server=None, **kwargs):
         self.__id = _next_id()
@@ -232,7 +250,6 @@ def clear_watchers(self):
 
     def clone(self):
         other = self.__class__(trame_server=self.server)
-        print(other)
         state = getattr(self, "_server_state", {})
         other.update(**state)
         return other
@@ -362,18 +379,38 @@ def decode_dataclass_item(item):
 def encode_dataclass_list(items):
     if items is None:
         return None
-    print("encode list", items)
+    # print("encode list", items)
     return [item._id for item in items]
 
 
 def decode_dataclass_list(items):
     # print("decode_dataclass_list", items)
     if items is None:
         return None
-    print("decode list", items)
+    # print("decode list", items)
     return list(map(get_instance, items))
 
 
+def decode_dataclass_set(items):
+    # print("decode_dataclass_list", items)
+    if items is None:
+        return None
+    # print("decode list", items)
+    return set(map(get_instance, items))
+
+
+def encode_set(items):
+    if items is None:
+        return None
+    return list(items)
+
+
+def decode_set(items):
+    if items is None:
+        return None
+    return set(items)
+
+
 def encode_dataclass_dict(data):
     if data is None:
         return None
@@ -396,6 +433,7 @@ def decode_dataclass_dict(data):
     "ServerOnly",
     "StateDataModel",
     "Sync",
+    "TypeValidation",
     "get_instance",
     "watch",
 ]
@@ -412,9 +450,9 @@ def __init__(
         self,
         _type,
         default=None,
-        convert=None,
-        has_dataclass=False,
-        type_checking="warning",  # error, warning, ignore
+        convert: FieldEncoder = None,
+        has_dataclass: bool = False,
+        type_checking: TypeValidation = TypeValidation.WARNING,
     ):
         self._type_checking = type_checking
         self._type = get_origin(_type) or _type
@@ -433,6 +471,9 @@ def __init__(
             if self._type is list:
                 encoder = encode_dataclass_list
                 decoder = decode_dataclass_list
+            elif self._type is set:
+                encoder = encode_dataclass_list
+                decoder = decode_dataclass_set
             elif self._type is dict:
                 encoder = encode_dataclass_dict
                 decoder = decode_dataclass_dict
@@ -442,13 +483,11 @@ def __init__(
 
             self._convert = FieldEncoder(encoder, decoder)
 
-    def __set_name__(self, owner, name):
-        if not hasattr(owner, "FIELD_NAMES"):
-            owner.FIELD_NAMES = set()
-
-        if not hasattr(owner, "TYPE_CHECKING"):
-            owner.TYPE_CHECKING = {}
+        if not self._convert and self._type is set:
+            self._convert = FieldEncoder(encode_set, decode_set)
 
+    def __set_name__(self, owner, name):
+        _setup_class_fields(owner)
         self._name = name
         owner.TYPE_CHECKING[name] = self._type_checking
         owner.FIELD_NAMES.add(name)
@@ -461,37 +500,25 @@ def __get__(self, instance, owner):
     def __set__(self, instance, value):
         type_check = instance.TYPE_CHECKING[self._name]
         if (
-            type_check in {"error", "warning"}
+            type_check in {TypeValidation.STRICT, TypeValidation.WARNING}
             and value is not None
             and not isinstance(value, self._type)
         ):
-            msg = f"{self._name} must be {self._type} instead of {type(value)}"
-            if type_check == "error":
+            msg = f"{self._name} must be {self._type} instead of {type(value)} for class {instance.__class__}"
+            if type_check == TypeValidation.STRICT:
                 raise TypeError(msg)
 
             logger.warning(msg)
 
-        instance._dirty_set.add(self._name)
-        instance._server_state[self._name] = value
-        instance._on_dirty()
+        if instance._server_state.get(self._name) != value:
+            instance._dirty_set.add(self._name)
+            instance._server_state[self._name] = value
+            instance._on_dirty()
 
 
 class Sync(ServerOnly):
     def __set_name__(self, owner, name):
-        if not hasattr(owner, "FIELD_NAMES"):
-            owner.FIELD_NAMES = set()
-
-        if not hasattr(owner, "TYPE_CHECKING"):
-            owner.TYPE_CHECKING = {}
-
-        if not hasattr(owner, "CLIENT_NAMES"):
-            owner.CLIENT_NAMES = set()
-
-        if not hasattr(owner, "ENCODERS"):
-            owner.ENCODERS = {}
-
-        if not hasattr(owner, "DATACLASS_NAMES"):
-            owner.DATACLASS_NAMES = set()
+        _setup_class_fields(owner)
 
         if self._has_dataclass:
             owner.DATACLASS_NAMES.add(name)
@@ -507,18 +534,7 @@ def __set_name__(self, owner, name):
 
 class ClientOnly(ServerOnly):
     def __set_name__(self, owner, name):
-        if not hasattr(owner, "FIELD_NAMES"):
-            owner.FIELD_NAMES = set()
-
-        if not hasattr(owner, "TYPE_CHECKING"):
-            owner.TYPE_CHECKING = {}
-
-        if not hasattr(owner, "CLIENT_NAMES"):
-            owner.CLIENT_NAMES = set()
-
-        if not hasattr(owner, "CLIENT_ONLY_NAMES"):
-            owner.CLIENT_ONLY_NAMES = set()
-
+        _setup_class_fields(owner)
         self._name = name
         owner.TYPE_CHECKING[name] = self._type_checking
         owner.FIELD_NAMES.add(name)

tests/test_dataclass.py

@@ -1,5 +1,4 @@
 import asyncio
-from contextlib import suppress
 from pathlib import Path
 
 import pytest
@@ -41,12 +40,12 @@ class Composite(StateDataModel):
 
 
 def test_complex_type():
-    with suppress(NonSerializableType):
+    with pytest.raises(NonSerializableType):
 
         class ErrorData(StateDataModel):
             path: Path
 
-        pytest.fail("Should trigger a NonSerializableType exception")
+        # pytest.fail("Should trigger a NonSerializableType exception")
 
 
 def test_watch():

tests/test_v2.py

@@ -0,0 +1,84 @@
+import pytest
+
+from trame_dataclass.v2 import (
+    ServerOnly,
+    StateDataModel,
+    Sync,
+    TypeValidation,
+)
+
+
+def test_input_validation():
+    class BasicTypeValidation(StateDataModel):
+        a = ServerOnly(int, 0, type_checking=TypeValidation.STRICT)
+        b = ServerOnly(float, 0.0, type_checking=TypeValidation.STRICT)
+        c = ServerOnly(str, "", type_checking=TypeValidation.STRICT)
+        d = ServerOnly(dict, dict, type_checking=TypeValidation.STRICT)
+        e = ServerOnly(list[int], list, type_checking=TypeValidation.STRICT)
+        f = ServerOnly(set, set, type_checking=TypeValidation.STRICT)
+        g = ServerOnly(str, "", type_checking=TypeValidation.SKIP)
+        h = ServerOnly(str, "", type_checking=TypeValidation.WARNING)
+
+    data = BasicTypeValidation()
+    data.a = 1
+    data.b = 1.1
+    data.c = "hello"
+    data.d = {"a": 1}
+    data.e = [1, 2, 3]
+    data.f = {"a", "b", "c"}
+    data.g = 5
+    data.h = 5
+
+    with pytest.raises(TypeError):
+        data.a = 1.2
+
+    with pytest.raises(TypeError):
+        data.b = 5
+
+    with pytest.raises(TypeError):
+        data.c = 0
+
+    with pytest.raises(TypeError):
+        data.d = []
+
+    with pytest.raises(TypeError):
+        data.e = {}
+
+    with pytest.raises(TypeError):
+        data.f = []
+
+
+def test_serialization():
+    class Dummy(StateDataModel):
+        a = Sync(int, 0)
+
+    class BasicSerial(StateDataModel):
+        a = Sync(int, 123)
+        b = Sync(Dummy, has_dataclass=True)
+        c = Sync(list[Dummy], list, has_dataclass=True)
+        d = Sync(dict[str, Dummy], dict, has_dataclass=True)
+        e = Sync(set, set)
+        f = Sync(set, set, has_dataclass=True)
+
+    data = BasicSerial()
+    data.b = Dummy()
+    data.c = [Dummy(), Dummy()]
+    data.d = {"a": Dummy(), "b": Dummy()}
+    data.e = {"a", "b", "c"}
+    data.f = {Dummy(), Dummy(), Dummy()}
+
+    state = data.client_state
+
+    assert state.get("a") == 123
+    assert state.get("b") == data.b._id
+    assert state.get("c") == [i._id for i in data.c]
+    assert state.get("d") == {k: v._id for k, v in data.d.items()}
+
+    set_as_list = state.get("e")
+    assert isinstance(set_as_list, list)
+    assert set(set_as_list) == {"a", "b", "c"}
+
+    set_as_list = state.get("f")
+    assert isinstance(set_as_list, list)
+    for obj in data.f:
+        assert obj._id in set_as_list