diff --git a/autogen/cache/redis_cache.py b/autogen/cache/redis_cache.py index 88f7a36b9a2..37bb8b00504 100644 --- a/autogen/cache/redis_cache.py +++ b/autogen/cache/redis_cache.py @@ -1,7 +1,15 @@ import pickle +from types import TracebackType +from typing import Any, Optional, Type import redis +import sys from .abstract_cache_base import AbstractCache +if sys.version_info >= (3, 11): + from typing import Self +else: + from typing_extensions import Self + class RedisCache(AbstractCache): """ @@ -24,7 +32,7 @@ class RedisCache(AbstractCache): __exit__(self, exc_type, exc_value, traceback): Context management exit. """ - def __init__(self, seed, redis_url): + def __init__(self, seed: str, redis_url: str): """ Initialize the RedisCache instance. @@ -36,7 +44,7 @@ def __init__(self, seed, redis_url): self.seed = seed self.cache = redis.Redis.from_url(redis_url) - def _prefixed_key(self, key): + def _prefixed_key(self, key: str) -> str: """ Get a namespaced key for the cache. @@ -48,7 +56,7 @@ def _prefixed_key(self, key): """ return f"autogen:{self.seed}:{key}" - def get(self, key, default=None): + def get(self, key: str, default: Optional[Any] = None) -> Any: """ Retrieve an item from the Redis cache. @@ -65,7 +73,7 @@ def get(self, key, default=None): return default return pickle.loads(result) - def set(self, key, value): + def set(self, key: str, value: Any) -> None: """ Set an item in the Redis cache. @@ -79,7 +87,7 @@ def set(self, key, value): serialized_value = pickle.dumps(value) self.cache.set(self._prefixed_key(key), serialized_value) - def close(self): + def close(self) -> None: """ Close the Redis client. @@ -87,7 +95,7 @@ def close(self): """ self.cache.close() - def __enter__(self): + def __enter__(self) -> Self: """ Enter the runtime context related to the object. @@ -96,7 +104,9 @@ def __enter__(self): """ return self - def __exit__(self, exc_type, exc_value, traceback): + def __exit__( + self, exc_type: Optional[Type[BaseException]], exc_val: Optional[BaseException], exc_tb: Optional[TracebackType] + ) -> None: """ Exit the runtime context related to the object.