Support more types for freezing

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)
+        args = tuple(_freeze(arg) for arg in args)
-        kwargs = {
+        kwargs = {k: _freeze(v) for k, v in kwargs.items()}
-            k: FrozenDict(v) if isinstance(v, dict) else 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