diff --git a/pyproject.toml b/pyproject.toml index 316a6ba..1f0cf07 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -31,6 +31,7 @@ dev = [ "ruff", "pytest", "pytest-asyncio", + "pytest-benchmark", "pytest-cov", "nox", "trame", diff --git a/src/trame_server/utils/typed_state.py b/src/trame_server/utils/typed_state.py index 5d9feed..ad58711 100644 --- a/src/trame_server/utils/typed_state.py +++ b/src/trame_server/utils/typed_state.py @@ -5,6 +5,8 @@ from datetime import date, datetime, time, timezone from decimal import Decimal from enum import Enum +from functools import cache, reduce +from operator import or_ from pathlib import Path from types import UnionType from typing import ( @@ -28,6 +30,52 @@ V = TypeVar("V") +def _resolve_type(annotation, type_var_mapping): + if isinstance(annotation, TypeVar): + return type_var_mapping.get(annotation, annotation) + + origin = get_origin(annotation) + if origin is None: + return annotation + + args = tuple( + _resolve_type(argument, type_var_mapping) for argument in get_args(annotation) + ) + if origin is UnionType: + return reduce(or_, args) + return origin[args] + + +@cache +def _get_resolved_type_hints(dataclass_type): + type_var_mapping = {} + + def visit(kls): + for base in getattr(kls, "__orig_bases__", ()): + origin = get_origin(base) + if origin is None: + continue + + arguments = get_args(base) + parameters = getattr(origin, "__parameters__", ()) + type_var_mapping.update( + zip( + parameters, + ( + _resolve_type(argument, type_var_mapping) + for argument in arguments + ), + ) + ) + visit(origin) + + visit(dataclass_type) + return { + name: _resolve_type(annotation, type_var_mapping) + for name, annotation in get_type_hints(dataclass_type).items() + } + + class _SerializationFailure: """ Simple class to handle encoding / decoding failures @@ -115,16 +163,21 @@ class CollectionEncoderDecoder(IStateEncoderDecoder): :param encoders: List of encoders to use when encoding/decoding lists and dicts. """ + _DATACLASS_TYPE_KEY = "__trame_typed_state__" + def __init__(self, encoders: list[IStateEncoderDecoder] | None = None): self._encoders = encoders or [DefaultEncoderDecoder()] def encode(self, obj): if is_dataclass(obj): - return { + value = { field.name: self.encode(getattr(obj, field.name)) for field in fields(obj) } + value[self._DATACLASS_TYPE_KEY] = self._get_dataclass_type_name(type(obj)) + return value + if isinstance(obj, dict): return {self.encode(key): self.encode(value) for key, value in obj.items()} @@ -158,20 +211,33 @@ def decode(self, obj, obj_type: type): raise TypeError(_error_msg) def _try_decode(self, obj, obj_type: type): - for decode in self._decode_strategies(): - val = decode(obj, obj_type) - if self.is_serialization_success(val): - return val - return self.failed_serialization() + if obj is None: + return None - def _decode_strategies(self) -> list[Callable[[Any, type], Any]]: - return [ - self._decode_dataclass, - self._decode_union, - self._decode_dict, - self._decode_iterable, - self._delegate_decode, - ] + origin = get_origin(obj_type) + if isinstance(obj, dict): + if ( + self._DATACLASS_TYPE_KEY in obj + or is_dataclass(obj_type) + or self._is_union_type(obj_type) + ): + return self._decode_dataclass_or_union(obj, obj_type) + if origin is dict: + return self._decode_dict(obj, obj_type) + elif self._is_union_type(obj_type): + return self._decode_union(obj, obj_type) + elif origin in (list, tuple) and self._is_iterable(obj): + return self._decode_iterable(obj, obj_type) + + return self._delegate_decode(obj, obj_type) + + def _decode_dataclass_or_union(self, obj, obj_type): + val = self._decode_dataclass(obj, obj_type) + if self.is_serialization_success(val): + return val + if self._is_union_type(obj_type): + return self._decode_union(obj, obj_type) + return self._delegate_decode(obj, obj_type) def _delegate_decode(self, obj, obj_type: type): for encoder in self._encoders: @@ -181,7 +247,7 @@ def _delegate_decode(self, obj, obj_type: type): return self.failed_serialization() def _decode_dict(self, obj, obj_type: type): - if not isinstance(obj, dict): + if not isinstance(obj, dict) or get_origin(obj_type) is not dict: return self.failed_serialization() key_type, value_type = get_args(obj_type) @@ -191,7 +257,7 @@ def _decode_dict(self, obj, obj_type: type): } def _decode_iterable(self, obj, obj_type: type): - if not self._is_iterable(obj): + if not self._is_iterable(obj) or get_origin(obj_type) not in (list, tuple): return self.failed_serialization() value_type = get_args(obj_type)[0] @@ -208,16 +274,43 @@ def _decode_union(self, obj, obj_type: type): return self.failed_serialization() def _decode_dataclass(self, obj, obj_type: type): + if not isinstance(obj, dict): + return self.failed_serialization() + + if self._DATACLASS_TYPE_KEY in obj: + obj_type = self._get_dataclass_type(obj, obj_type) + if not is_dataclass(obj_type): return self.failed_serialization() - field_types = get_type_hints(obj_type) + field_types = _get_resolved_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) + @staticmethod + def _get_dataclass_type_name(dataclass_type: type) -> str: + return f"{dataclass_type.__module__}.{dataclass_type.__qualname__}" + + def _get_dataclass_type(self, obj, obj_type: type) -> type | None: + type_name = obj[self._DATACLASS_TYPE_KEY] + dataclass_types = ( + (obj_type,) + if is_dataclass(obj_type) + else get_args(obj_type) + if self._is_union_type(obj_type) + else () + ) + for dataclass_type in dataclass_types: + if ( + is_dataclass(dataclass_type) + and self._get_dataclass_type_name(dataclass_type) == type_name + ): + return dataclass_type + return None + @classmethod def _is_union_type(cls, obj_type: type): return get_origin(obj_type) is Union or isinstance(obj_type, UnionType) @@ -261,7 +354,7 @@ def __init__( if default_value == MISSING and default_factory != MISSING: default_value = default_factory() if default_value != MISSING: - self._state.setdefault(self._state_id, self._encoder.encode(default_value)) + self.set_default_value(default_value) def __get__(self, instance, owner): return self.get_value() @@ -274,7 +367,40 @@ def get_value(self): return self._encoder.decode(value, self._type) def set_value(self, value): - self._state[self._state_id] = self._encoder.encode(value) + self._state[self._state_id] = self._encode(value) + + def set_default_value(self, value): + self._state.setdefault(self._state_id, self._encode(value)) + + def _encode(self, value): + return self._encoder.encode(value) + + +class _ProxyDataclassField: + """ + Descriptor for nested dataclass state proxies. + """ + + def __init__(self, proxy, default, default_factory): + self._proxy = proxy + + default_value = default + if default_value == MISSING and default_factory != MISSING: + default_value = default_factory() + if default_value != MISSING: + self.set_default_value(default_value) + + def __get__(self, instance, owner): + return self._proxy + + def __set__(self, instance, value): + TypedState.from_dataclass(self._proxy, value) + + def set_default_value(self, dataclass_obj: T) -> None: + for f in fields(TypedState._get_proxy_dataclass_type_or_raise(self._proxy)): + value = getattr(dataclass_obj, f.name) + proxy_field = vars(type(self._proxy))[f.name] + proxy_field.set_default_value(value) class _NameField: @@ -305,6 +431,7 @@ class TypedState(Generic[T]): _STATE_PROXY_DATACLASS_TYPE = "__state_proxy_dataclass_type" _STATE_PROXY_FIELD_DICT = "__state_proxy_field_dict" _STATE_PROXY_STATE_ID = "__state_proxy_state_id" + _DATA_PROXY_CLASS_SUFFIX = "__Proxy" def __init__( self, @@ -399,7 +526,9 @@ def handler(state_id: str, field: Field, field_type: type): state_encoder=encoder, ) - return cls._build_proxy_cls(dataclass_type, namespace, handler, "__Proxy") + return cls._build_proxy_cls( + dataclass_type, namespace, handler, cls._DATA_PROXY_CLASS_SUFFIX + ) @classmethod def _create_state_names_proxy(cls, dataclass_type: Type[T], *, namespace="") -> T: @@ -445,18 +574,25 @@ def _build_proxy_cls( # Use type hints instead of field.type to avoid lazy evaluation of field.type when used in files containing # from __future__ import annotations header. - field_types = get_type_hints(dataclass_type) + field_types = _get_resolved_type_hints(dataclass_type) for f in fields(dataclass_type): state_id = f"{prefix}__{f.name}" f_type = field_types[f.name] if is_dataclass(f_type): - field = cls._build_proxy_cls( + nested_proxy = cls._build_proxy_cls( f_type, state_id, handler, cls_suffix, inner_field_dict ) + field = ( + _ProxyDataclassField(nested_proxy, f.default, f.default_factory) + if cls_suffix == cls._DATA_PROXY_CLASS_SUFFIX + else nested_proxy + ) + inner_field_dict[cls.get_state_id(nested_proxy, state_id)] = ( + nested_proxy + ) else: field = handler(state_id, f, f_type) - - inner_field_dict[cls.get_state_id(field, state_id)] = field + inner_field_dict[cls.get_state_id(field, state_id)] = field namespace[f.name] = field if proxy_field_dict is not None: @@ -510,7 +646,7 @@ def is_name_proxy_class(cls, instance: T) -> bool: @classmethod def is_data_proxy_class(cls, instance: T) -> bool: return cls.is_proxy_class(instance) and type(instance).__name__.endswith( - "__Proxy" + cls._DATA_PROXY_CLASS_SUFFIX ) @classmethod @@ -550,11 +686,11 @@ def from_dataclass(cls, instance: T, dataclass_obj: T) -> None: raise TypeError(_error_msg) for f in fields(dataclass_type): - attr = getattr(instance, f.name) value = getattr(dataclass_obj, f.name) + proxy_field = vars(type(instance))[f.name] - if cls.is_proxy_class(attr): - cls.from_dataclass(attr, value) + if isinstance(proxy_field, _ProxyDataclassField): + cls.from_dataclass(proxy_field._proxy, value) else: setattr(instance, f.name, value) diff --git a/tests/test_typed_state.py b/tests/test_typed_state.py index 6bf816d..0b6e71d 100644 --- a/tests/test_typed_state.py +++ b/tests/test_typed_state.py @@ -3,6 +3,7 @@ from datetime import date, datetime, time, timezone from enum import Enum, auto from pathlib import Path +from typing import Generic, TypeVar from unittest.mock import MagicMock from uuid import UUID, uuid4 @@ -10,11 +11,14 @@ from trame_server import Server from trame_server.utils.typed_state import ( + CollectionEncoderDecoder, DefaultEncoderDecoder, IStateEncoderDecoder, TypedState, ) +V = TypeVar("V") + @pytest.fixture def server(): @@ -46,6 +50,38 @@ class MyBiggerData: c: float = 42.0 +@dataclass +class GenericItem: + value: int + + +@dataclass +class GenericState(Generic[V]): + current_index: int | None = None + items: list[V] = field(default_factory=list) + + +@dataclass +class GenericItemState(GenericState[GenericItem]): + pass + + +@dataclass +class MyTypeA: + value_a: int + + +@dataclass +class MyTypeB: + value_b: str + + +@dataclass +class UnionState: + current_index: int | None = None + items: list[MyTypeA | MyTypeB] = field(default_factory=list) + + def test_can_be_constructed_from_simple_dataclass(state): typed_state = TypedState(state, MyData) @@ -62,6 +98,56 @@ def test_can_be_constructed_from_nested_dataclass(state): assert state[typed_state.name.my_other_data.a] == 42 +def test_resolves_inherited_generic_dataclass_field_types(state): + typed_state = TypedState(state, GenericItemState) + + typed_state.data.current_index = 0 + typed_state.data.items = [GenericItem(value=42)] + + assert typed_state.data.current_index == 0 + _type_key = CollectionEncoderDecoder._DATACLASS_TYPE_KEY + assert state[typed_state.name.items] == [ + {_type_key: "test_typed_state.GenericItem", "value": 42} + ] + assert typed_state.data.items == [GenericItem(value=42)] + assert typed_state.get_dataclass() == GenericItemState( + current_index=0, items=[GenericItem(value=42)] + ) + + +def test_can_deserializes_non_union_dataclass_without_type_discriminator(state): + typed_state = TypedState(state, GenericItemState) + + state[typed_state.name.items] = [{"value": 42}] + + assert typed_state.data.items == [GenericItem(value=42)] + assert typed_state.get_dataclass() == GenericItemState( + items=[GenericItem(value=42)] + ) + + +def test_can_serializes_union_dataclass_type_in_collection(state): + typed_state = TypedState(state, UnionState) + expected_items = [ + MyTypeB(value_b="first"), + MyTypeA(value_a=42), + MyTypeA(value_a=43), + MyTypeB(value_b="last"), + ] + + typed_state.data.items = expected_items + + _type_key = CollectionEncoderDecoder._DATACLASS_TYPE_KEY + assert state[typed_state.name.items] == [ + {_type_key: "test_typed_state.MyTypeB", "value_b": "first"}, + {_type_key: "test_typed_state.MyTypeA", "value_a": 42}, + {_type_key: "test_typed_state.MyTypeA", "value_a": 43}, + {_type_key: "test_typed_state.MyTypeB", "value_b": "last"}, + ] + assert typed_state.data.items == expected_items + assert typed_state.get_dataclass() == UnionState(items=expected_items) + + def test_can_handle_namespace(state): typed_state_ns1 = TypedState(state, MyData, namespace="ns1") typed_state_ns2 = TypedState(state, MyData, namespace="ns2") @@ -92,9 +178,14 @@ def test_can_be_set_from_data_class(state): assert typed_state.data.my_other_data.b == 43 assert typed_state.data.c == 44 - TypedState.from_dataclass(typed_state.data.my_other_data, MyData(a=45, b=46)) + typed_state.data.my_other_data = MyData(a=45, b=46) assert typed_state.data.my_other_data.a == 45 assert typed_state.data.my_other_data.b == 46 + assert state[typed_state.name.my_other_data.a] == 45 + assert state[typed_state.name.my_other_data.b] == 46 + assert typed_state.get_dataclass() == MyBiggerData( + my_other_data=MyData(a=45, b=46), c=44 + ) def test_can_be_used_to_connect_to_state_changes(state): @@ -397,6 +488,21 @@ class DataWithUnionTypes: my_optional_enum_list: list[MyEnum | str | None] | None +@dataclass +class DataWithFloatOrFloatList: + t: float | list[float] + + +def test_deserializes_float_or_float_list(state): + typed_state = TypedState(state, DataWithFloatOrFloatList) + + state[typed_state.name.t] = 42.5 + assert typed_state.data.t == 42.5 + + state[typed_state.name.t] = [1.5, 2.5] + assert typed_state.data.t == [1.5, 2.5] + + def test_handles_union_types_by_order_of_definition(state): encode = CustomEnumEncode() typed_state = TypedState(state, DataWithUnionTypes, encoders=[encode]) @@ -468,6 +574,26 @@ class SimpleTypes: my_path: Path +@dataclass +class NonDefaultInit: + simple_types: SimpleTypes = field( + default_factory=lambda: SimpleTypes( + my_int=42, my_enum=MyEnum.B, my_path=Path("./shared") + ) + ) + + +def test_initializes_nested_dataclass_from_default_factory(state): + typed_state = TypedState(state, NonDefaultInit) + + assert typed_state.data.simple_types.my_int == 42 + assert typed_state.data.simple_types.my_enum == MyEnum.B + assert typed_state.data.simple_types.my_path == Path("shared") + assert state[typed_state.name.simple_types.my_int] == 42 + assert state[typed_state.name.simple_types.my_enum] == MyEnum.B.value + assert state[typed_state.name.simple_types.my_path] == "shared" + + @dataclass class TypedComposite: simple_types: SimpleTypes = field(default_factory=SimpleTypes) @@ -479,6 +605,43 @@ class DataclassCollections: nested_dict: dict[str, TypedComposite] = field(default_factory=dict) +@dataclass +class BenchmarkState(Generic[V]): + collections: DataclassCollections + selected: V | None = None + + +@dataclass +class BenchmarkTypedState(BenchmarkState[TypedComposite]): + pass + + +@pytest.fixture +def benchmark_data(): + composites = [ + TypedComposite( + SimpleTypes( + my_int=index, + my_enum=MyEnum(index % 3 + 1), + my_path=Path(f"/path/to/{index}"), + ) + ) + for index in range(42) + ] + collections = DataclassCollections( + nested_list=composites, + nested_dict={f"item-{index}": item for index, item in enumerate(composites)}, + ) + return BenchmarkTypedState(collections=collections, selected=composites[0]) + + +@pytest.fixture +def benchmark_typed_state(state, benchmark_data): + typed_state = TypedState(state, BenchmarkTypedState) + typed_state.set_dataclass(benchmark_data) + return typed_state + + def test_encode_decode_supports_collections_of_nested_dataclass(state): typed_state = TypedState(state, DataclassCollections) typed_state.data.nested_list = [ @@ -512,20 +675,25 @@ def test_encode_decode_supports_collections_of_nested_dataclass(state): SimpleTypes(my_int=4, my_enum=MyEnum.A, my_path=Path("/path/to/4")) ) + _type_key = CollectionEncoderDecoder._DATACLASS_TYPE_KEY 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"), - } + _type_key: "test_typed_state.SimpleTypes", + }, + _type_key: "test_typed_state.TypedComposite", }, { "simple_types": { "my_int": 2, "my_enum": typed_state.encode(MyEnum.B), "my_path": typed_state.encode("/path/to/2"), - } + _type_key: "test_typed_state.SimpleTypes", + }, + _type_key: "test_typed_state.TypedComposite", }, ] @@ -535,14 +703,18 @@ def test_encode_decode_supports_collections_of_nested_dataclass(state): "my_int": 3, "my_enum": typed_state.encode(MyEnum.C), "my_path": typed_state.encode("/path/to/3"), - } + _type_key: "test_typed_state.SimpleTypes", + }, + _type_key: "test_typed_state.TypedComposite", }, "4": { "simple_types": { "my_int": 4, "my_enum": typed_state.encode(MyEnum.A), "my_path": typed_state.encode("/path/to/4"), - } + _type_key: "test_typed_state.SimpleTypes", + }, + _type_key: "test_typed_state.TypedComposite", }, } @@ -563,3 +735,13 @@ def test_can_supress_change_listeners(state): state.flush() mock.assert_not_called() + + +def test_benchmark_typed_state_serialization( + benchmark, benchmark_typed_state, benchmark_data +): + benchmark(benchmark_typed_state.set_dataclass, benchmark_data) + + +def test_benchmark_typed_state_deserialization(benchmark, benchmark_typed_state): + benchmark(benchmark_typed_state.get_dataclass)