Improve typing hints in valkey/connection.py

Signed-off-by: Mikhail Koviazin <mikhail.koviazin@aiven.io>

Author: Mikhail Koviazin

valkey/_cache.py

@@ -9,7 +9,7 @@
 from valkey.typing import KeyT, ResponseT
 
 
-class EvictionPolicy(Enum):
+class EvictionPolicy(str, Enum):
     LRU = "lru"
     LFU = "lfu"
     RANDOM = "random"

valkey/connection.py

@@ -23,7 +23,7 @@
 from os import getpid
 from queue import Empty, Full, LifoQueue
 from time import time
-from typing import Any, Callable, List, Optional, Sequence, Type, Union
+from typing import Any, Callable, Iterable, List, Optional, Sequence, Tuple, Type, Union
 
 from ._cache import (
     DEFAULT_ALLOW_LIST,
@@ -46,7 +46,7 @@
     ValkeyError,
 )
 from .retry import Retry
-from .typing import KeysT, ResponseT
+from .typing import KeyT, ResponseT
 from .utils import (
     CRYPTOGRAPHY_AVAILABLE,
     LIBVALKEY_AVAILABLE,
@@ -66,8 +66,6 @@
 
 DEFAULT_RESP_VERSION = 2
 
-SENTINEL = object()
-
 DefaultParser: Type[Union[_RESP2Parser, _RESP3Parser, _LibvalkeyParser]]
 if LIBVALKEY_AVAILABLE:
     DefaultParser = _LibvalkeyParser
@@ -76,9 +74,9 @@
 
 
 class LibvalkeyRespSerializer:
-    def pack(self, *args: List):
+    def pack(self, *args) -> List[bytes]:
         """Pack a series of arguments into the Valkey protocol"""
-        output = []
+        output: List[bytes] = []
 
         if isinstance(args[0], str):
             args = tuple(args[0].encode().split()) + args[1:]
@@ -98,7 +96,7 @@ def __init__(self, buffer_cutoff, encode) -> None:
         self._buffer_cutoff = buffer_cutoff
         self.encode = encode
 
-    def pack(self, *args):
+    def pack(self, *args) -> List[bytes]:
         """Pack a series of arguments into the Valkey protocol"""
         output = []
         # the client might have included 1 or more literal arguments in
@@ -154,7 +152,7 @@ def __init__(
         socket_timeout: Optional[float] = 5,
         socket_connect_timeout: Optional[float] = None,
         retry_on_timeout: bool = False,
-        retry_on_error=SENTINEL,
+        retry_on_error: Optional[List[Type[Exception]]] = None,
         encoding: str = "utf-8",
         encoding_errors: str = "strict",
         decode_responses: bool = False,
@@ -166,15 +164,15 @@ def __init__(
         lib_version: Optional[str] = get_lib_version(),
         username: Optional[str] = None,
         retry: Union[Any, None] = None,
-        valkey_connect_func: Optional[Callable[[], None]] = None,
+        valkey_connect_func: Optional[Callable[["AbstractConnection"], None]] = None,
         credential_provider: Optional[CredentialProvider] = None,
-        protocol: Optional[int] = 2,
+        protocol: Optional[Union[int, str]] = 2,
         command_packer: Optional[Callable[[], None]] = None,
         cache_enabled: bool = False,
         client_cache: Optional[AbstractCache] = None,
         cache_max_size: int = 10000,
         cache_ttl: int = 0,
-        cache_policy: str = DEFAULT_EVICTION_POLICY,
+        cache_policy=DEFAULT_EVICTION_POLICY,
         cache_deny_list: List[str] = DEFAULT_DENY_LIST,
         cache_allow_list: List[str] = DEFAULT_ALLOW_LIST,
     ):
@@ -205,7 +203,7 @@ def __init__(
             socket_connect_timeout = socket_timeout
         self.socket_connect_timeout = socket_connect_timeout
         self.retry_on_timeout = retry_on_timeout
-        if retry_on_error is SENTINEL:
+        if retry_on_error is None:
             retry_on_error = []
         if retry_on_timeout:
             # Add TimeoutError to the errors list to retry on
@@ -222,18 +220,16 @@ def __init__(
         else:
             self.retry = Retry(NoBackoff(), 0)
         self.health_check_interval = health_check_interval
-        self.next_health_check = 0
+        self.next_health_check = 0.0
         self.valkey_connect_func = valkey_connect_func
         self.encoder = Encoder(encoding, encoding_errors, decode_responses)
-        self._sock = None
+        self._sock: Optional[socket.socket] = None
         self._socket_read_size = socket_read_size
         self.set_parser(parser_class)
-        self._connect_callbacks = []
+        self._connect_callbacks: List[weakref.WeakMethod] = []
         self._buffer_cutoff = 6000
         try:
-            p = int(protocol)
-        except TypeError:
-            p = DEFAULT_RESP_VERSION
+            p = int(protocol) if protocol is not None else DEFAULT_RESP_VERSION
         except ValueError:
             raise ConnectionError("protocol must be an integer")
         finally:
@@ -260,7 +256,7 @@ def __repr__(self):
         return f"<{self.__class__.__module__}.{self.__class__.__name__}({repr_args})>"
 
     @abstractmethod
-    def repr_pieces(self):
+    def repr_pieces(self) -> List[Tuple[str, Any]]:
         pass
 
     def __del__(self):
@@ -310,7 +306,7 @@ def set_parser(self, parser_class):
 
     def connect(self):
         "Connects to the Valkey server if not already connected"
-        if self._sock:
+        if self._sock is not None:
             return
         try:
             sock = self.retry.call_with_retry(
@@ -348,7 +344,7 @@ def _connect(self):
         pass
 
     @abstractmethod
-    def _host_error(self):
+    def _host_error(self) -> str:
         pass
 
     def _error_message(self, exception):
@@ -377,7 +373,7 @@ def on_connect(self):
                 self._parser.EXCEPTION_CLASSES = parser.EXCEPTION_CLASSES
                 self._parser.on_connect(self)
             if len(auth_args) == 1:
-                auth_args = ["default", auth_args[0]]
+                auth_args = ("default", auth_args[0])
             self.send_command("HELLO", self.protocol, "AUTH", *auth_args)
             response = self.read_response()
             # if response.get(b"proto") != self.protocol and response.get(
@@ -492,6 +488,7 @@ def send_packed_command(self, command, check_health=True):
         """Send an already packed command to the Valkey server"""
         if not self._sock:
             self.connect()
+        assert self._sock is not None
         # guard against health check recursion
         if check_health:
             self.check_health()
@@ -592,8 +589,8 @@ def pack_command(self, *args):
 
     def pack_commands(self, commands):
         """Pack multiple commands into the Valkey protocol"""
-        output = []
-        pieces = []
+        output: List[bytes] = []
+        pieces: List[bytes] = []
         buffer_length = 0
         buffer_cutoff = self._buffer_cutoff
 
@@ -629,6 +626,7 @@ def _cache_invalidation_process(
         and the second string is the list of keys to invalidate.
         (if the list of keys is None, then all keys are invalidated)
         """
+        assert self.client_cache is not None
         if data[1] is None:
             self.client_cache.flush()
         else:
@@ -650,7 +648,7 @@ def _get_from_local_cache(self, command: Sequence[str]):
         return self.client_cache.get(command)
 
     def _add_to_local_cache(
-        self, command: Sequence[str], response: ResponseT, keys: List[KeysT]
+        self, command: Sequence[str], response: ResponseT, keys: List[KeyT]
     ):
         """
         Add the command and response to the local cache if the command
@@ -671,7 +669,7 @@ def delete_command_from_cache(self, command: Union[str, Sequence[str]]):
         if self.client_cache:
             self.client_cache.delete_command(command)
 
-    def invalidate_key_from_cache(self, key: KeysT):
+    def invalidate_key_from_cache(self, key: KeyT):
         if self.client_cache:
             self.client_cache.invalidate_key(key)
 
@@ -1036,7 +1034,7 @@ def __init__(
         self._fork_lock = threading.Lock()
         self.reset()
 
-    def __repr__(self) -> (str, str):
+    def __repr__(self) -> str:
         return (
             f"<{type(self).__module__}.{type(self).__name__}"
             f"({repr(self.connection_class(**self.connection_kwargs))})>"
@@ -1045,8 +1043,8 @@ def __repr__(self) -> (str, str):
     def reset(self) -> None:
         self._lock = threading.Lock()
         self._created_connections = 0
-        self._available_connections = []
-        self._in_use_connections = set()
+        self._available_connections: list[Connection] = []
+        self._in_use_connections: set[Connection] = set()
 
         # this must be the last operation in this method. while reset() is
         # called when holding _fork_lock, other threads in this process
@@ -1193,7 +1191,7 @@ def disconnect(self, inuse_connections: bool = True) -> None:
         self._checkpid()
         with self._lock:
             if inuse_connections:
-                connections = chain(
+                connections: Iterable[Connection] = chain(
                     self._available_connections, self._in_use_connections
                 )
             else:
@@ -1388,7 +1386,7 @@ def release(self, connection):
             # we don't want this connection
             pass
 
-    def disconnect(self):
+    def disconnect(self, _: bool = True) -> None:
         "Disconnects all connections in the pool."
         self._checkpid()
         for connection in self._connections:

valkey/retry.py

@@ -7,7 +7,7 @@
 T = TypeVar("T")
 
 if TYPE_CHECKING:
-    from redis.backoff import AbstractBackoff
+    from valkey.backoff import AbstractBackoff
 
 
 class Retry: