Support more types for freezing

This commit is contained in:
Andras Schmelczer 2022-07-01 19:50:26 +02:00
parent c8fbced859
commit 8f1418c346

View file

@ -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