diff --git a/src/great_ai/great_ai/helper/freeze_arguments.py b/src/great_ai/great_ai/helper/freeze_arguments.py index 3847ea5..92e5e51 100644 --- a/src/great_ai/great_ai/helper/freeze_arguments.py +++ b/src/great_ai/great_ai/helper/freeze_arguments.py @@ -1,12 +1,22 @@ from functools import wraps -from typing import Any, Callable, Dict, List +from typing import Any, Callable, Dict, List, Set, Union class FrozenDict(dict): - def __hash__(self) -> int: # type: ignore + def __hash__(self) -> int: return hash(frozenset(self.items())) +class FrozenList(list): + def __hash__(self) -> int: + return hash(tuple(self)) + + +class FrozenSet(set): + def __hash__(self) -> int: + return hash(frozenset(self)) + + def freeze_arguments(func: Callable[..., Any]) -> Callable[..., Any]: """Transform mutable dictionary Into immutable @@ -16,10 +26,21 @@ def freeze_arguments(func: Callable[..., Any]) -> Callable[..., Any]: @wraps(func) def wrapper(*args: List[Any], **kwargs: Dict[str, Any]) -> Any: - args = tuple(FrozenDict(arg) if isinstance(arg, dict) else arg for arg in args) - kwargs = { - k: FrozenDict(v) if isinstance(v, dict) else 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: + if isinstance(value, dict): + return FrozenDict(value) + + if isinstance(value, list): + return FrozenList(value) + + if isinstance(value, set): + return FrozenSet(value) + + return value