andras/great-ai
1
0
Fork 0
Code Issues Pull requests 1 Projects Releases 4 Packages Wiki Activity Actions

Serialise BaseModels

Browse source
This commit is contained in:
Andras Schmelczer 2022-07-03 16:04:27 +02:00
parent 2f4d203446
commit 623e81cdcb
3 changed files with 18 additions and 8 deletions
Download patch file Download diff file Expand all files Collapse all files

2
src/great_ai/great_ai/helper/__init__.py
View file

@ -1,4 +1,4 @@
from .freeze_arguments import freeze_arguments from .freeze_arguments import freeze, freeze_arguments
from .get_arguments import get_arguments from .get_arguments import get_arguments
from .get_function_metadata_store import get_function_metadata_store from .get_function_metadata_store import get_function_metadata_store
from .hashable_base_model import HashableBaseModel from .hashable_base_model import HashableBaseModel

22
src/great_ai/great_ai/helper/freeze_arguments.py
View file

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

2
src/great_ai/great_ai/persistence/mongodb_driver.py
View file

@ -121,7 +121,7 @@ class MongodbDriver(TracingDatabaseDriver):
return documents, len(documents) return documents, len(documents)
def update(self, id: str, new_version: Trace) -> None: def update(self, id: str, new_version: Trace) -> None:
serialized = new_version.dict() serialized = new_version.to_flat_dict()
serialized["_id"] = new_version.trace_id serialized["_id"] = new_version.trace_id
with MongoClient(self.mongo_connection_string) as client: with MongoClient(self.mongo_connection_string) as client:
client[self.mongo_database].traces.update_one(id, new_version) client[self.mongo_database].traces.update_one(id, new_version)