Serialise BaseModels

src/great_ai/great_ai/helper/freeze_arguments.py

@@ -1,20 +1,22 @@
 from functools import wraps
 from typing import Any, Callable, Dict, List, Set, Union
 
+from pydantic import BaseModel
+
 
 class FrozenDict(dict):
     def __hash__(self) -> int:
-        return hash(frozenset(self.items()))
+        return hash(frozenset((k, freeze(v)) for k, v in self.items()))
 
 
 class FrozenList(list):
     def __hash__(self) -> int:
-        return hash(tuple(self))
+        return hash(tuple(freeze(i) for i in self))
 
 
 class FrozenSet(set):
     def __hash__(self) -> int:
-        return hash(frozenset(self))
+        return hash(frozenset(freeze(i) for i in self))
 
 
 def freeze_arguments(func: Callable[..., Any]) -> Callable[..., Any]:
@@ -26,14 +28,14 @@ def freeze_arguments(func: Callable[..., Any]) -> Callable[..., Any]:
 
     @wraps(func)
     def wrapper(*args: List[Any], **kwargs: Dict[str, Any]) -> Any:
-        args = tuple(_freeze(arg) for arg in args)
-        kwargs = {k: _freeze(v) for k, v in kwargs.items()}
+        args = tuple(freeze(arg) for arg in args)
+        kwargs = {k: freeze(v) for k, v in kwargs.items()}
         return func(*args, **kwargs)
 
     return wrapper
 
 
-def _freeze(value: Union[List[Any], Dict[str, Any], Set[Any]]) -> Any:
+def freeze(value: Union[List[Any], Dict[str, Any], Set[Any]]) -> Any:
     if isinstance(value, dict):
         return FrozenDict(value)
 
@@ -43,4 +45,12 @@ def _freeze(value: Union[List[Any], Dict[str, Any], Set[Any]]) -> Any:
     if isinstance(value, set):
         return FrozenSet(value)
 
+    if isinstance(value, BaseModel):
+
+        class HashableValue(type(value)):
+            def __hash__(self) -> int:
+                return hash(frozenset((k, freeze(v)) for k, v in self.dict().items()))
+
+        return HashableValue(**value.dict())
+
     return value