Fix typing and minor issues

great_ai/helper/freeze_arguments.py

@@ -1,41 +1,50 @@
+import inspect
 from functools import wraps
-from typing import Any, Callable, Dict, List, Set, Union
+from typing import Any, Callable, Mapping, Sequence, Set, TypeVar, Union, cast
 
 from pydantic import BaseModel
 
+F = TypeVar("F", bound=Callable)
+
 
 class FrozenDict(dict):
-    def __hash__(self) -> int:
+    def __hash__(self) -> int:  # type: ignore
         return hash(frozenset((k, freeze(v)) for k, v in self.items()))
 
 
 class FrozenList(list):
-    def __hash__(self) -> int:
+    def __hash__(self) -> int:  # type: ignore
         return hash(tuple(freeze(i) for i in self))
 
 
 class FrozenSet(set):
-    def __hash__(self) -> int:
+    def __hash__(self) -> int:  # type: ignore
         return hash(frozenset(freeze(i) for i in self))
 
 
-def freeze_arguments(func: Callable[..., Any]) -> Callable[..., Any]:
-    """Transform mutable dictionary
-    Into immutable
-    Useful to be compatible with cache
-    source: https://stackoverflow.com/questions/6358481/using-functools-lru-cache-with-dictionary-arguments
-    """
-
+def freeze_arguments(func: F) -> F:
     @wraps(func)
-    def wrapper(*args: List[Any], **kwargs: Dict[str, Any]) -> Any:
+    def wrapper(*args: Any, **kwargs: Any) -> Any:
         args = tuple(freeze(arg) for arg in args)
         kwargs = {k: freeze(v) for k, v in kwargs.items()}
         return func(*args, **kwargs)
 
-    return wrapper
+    @wraps(func)
+    async def async_wrapper(*args: Any, **kwargs: Any) -> Any:
+        return await wrapper(*args, **kwargs)
+
+    return cast(F, async_wrapper if inspect.iscoroutinefunction(func) else wrapper)
 
 
-def freeze(value: Union[List[Any], Dict[str, Any], Set[Any]]) -> Any:
+def freeze(value: Union[Sequence[Any], Mapping[str, Any], Set[Any], BaseModel]) -> Any:
+    """
+    >>> class MyClass(BaseModel):
+    ...     a: int
+    >>> my_object = MyClass(a=3)
+    >>> my_other_object = MyClass(a=4)
+    >>> freeze(my_object) == freeze(my_other_object), freeze(my_object) == freeze(my_object)
+    (False, True)
+    """
     if isinstance(value, dict):
         return FrozenDict(value)
 
@@ -47,7 +56,7 @@ def freeze(value: Union[List[Any], Dict[str, Any], Set[Any]]) -> Any:
 
     if isinstance(value, BaseModel):
 
-        class HashableValue(type(value)):
+        class HashableValue(type(value)):  # type: ignore
             def __hash__(self) -> int:
                 return hash(frozenset((k, freeze(v)) for k, v in self.dict().items()))