Begin resolving type annotation issues

Author: kurtmckee

.github/workflows/integration.yaml

@@ -206,3 +206,21 @@ jobs:
           PIP_DISABLE_PIP_VERSION_CHECK: "1"
         run: |
           pip install --quiet git+${GITHUB_SERVER_URL}/${GITHUB_REPOSITORY}.git@${GITHUB_SHA}
+
+
+   test-type-annotations:
+     runs-on: ubuntu-latest
+     timeout-minutes: 5
+     name: Test type annotations
+     steps:
+       - uses: actions/checkout@v6
+       - uses: actions/setup-python@v6
+         with:
+           python-version: '3.14'
+           cache: 'pip'
+
+       - name: Test type annotations
+         run: |
+           pip install --group type-check
+           pip install .
+           mypy

.mypy.ini

@@ -56,6 +56,12 @@ dev = [
     'uvloop; implementation_name != "pypy"',
     'vulture>=2.3.0',
 ]
+type-check = [
+    'mypy',
+    'types-requests',
+    'cryptography',
+    'libvalkey',
+]
 docs = [
     'sphinx>=5.0,<7.0',
     'docutils<0.18',
@@ -103,3 +109,21 @@ markers = [
 ]
 asyncio_mode = "auto"
 timeout = 30
+
+[tool.mypy]
+packages = ["valkey"]
+sqlite_cache = true
+
+[[tool.mypy.overrides]]
+# The goal is to eliminate the `tool.mypy.overrides` section.
+# This will be accomplished by resolving type annotation issues
+# in the modules listed below.
+ignore_errors = true
+module = [
+    "valkey.client",
+    "valkey.cluster",
+    "valkey.connection",
+    "valkey.commands.*",
+    "valkey.asyncio.*",
+    "valkey._parsers.*",
+]

pyproject.toml

@@ -4,7 +4,7 @@
 from abc import ABC, abstractmethod
 from collections import OrderedDict, defaultdict
 from enum import Enum
-from typing import List, Sequence, Union
+from typing import List, Sequence, Union, Any
 
 from valkey.typing import KeyT, ResponseT
 
@@ -223,9 +223,9 @@ def __init__(
         self.max_size = max_size
         self.ttl = ttl
         self.eviction_policy = eviction_policy
-        self.cache = OrderedDict()
-        self.key_commands_map = defaultdict(set)
-        self.commands_ttl_list = []
+        self.cache: OrderedDict[str | Sequence[str], dict[str, Any]] = OrderedDict()
+        self.key_commands_map: defaultdict[KeyT, set] = defaultdict(set)
+        self.commands_ttl_list: list[str | Sequence[str]] = []
 
     def set(
         self,
@@ -252,7 +252,7 @@ def set(
         self._update_key_commands_map(keys_in_command, command)
         self.commands_ttl_list.append(command)
 
-    def get(self, command: Union[str, Sequence[str]]) -> ResponseT:
+    def get(self, command: Union[str, Sequence[str]]) -> ResponseT | None:
         """
         Get the response for a valkey command from the cache.
 
@@ -262,12 +262,14 @@ def get(self, command: Union[str, Sequence[str]]) -> ResponseT:
         Returns:
             ResponseT: The response associated with the command, or None if the command is not in the cache.  # noqa
         """
-        if command in self.cache:
-            if self._is_expired(command):
-                self.delete_command(command)
-                return
-            self._update_access(command)
-            return copy.deepcopy(self.cache[command]["response"])
+        if command not in self.cache:
+            return None
+
+        if self._is_expired(command):
+            self.delete_command(command)
+            return None
+        self._update_access(command)
+        return copy.deepcopy(self.cache[command]["response"])
 
     def delete_command(self, command: Union[str, Sequence[str]]):
         """
@@ -278,7 +280,7 @@ def delete_command(self, command: Union[str, Sequence[str]]):
         """
         if command in self.cache:
             keys_in_command = self.cache[command].get("keys")
-            self._del_key_commands_map(keys_in_command, command)
+            self._del_key_commands_map(keys_in_command, command)  # type: ignore[arg-type]
             self.commands_ttl_list.remove(command)
             del self.cache[command]
 

valkey/_cache.py

@@ -98,10 +98,10 @@ def __init__(self, cap: float = DEFAULT_CAP, base: float = DEFAULT_BASE) -> None
         """
         self._cap = cap
         self._base = base
-        self._previous_backoff = 0
+        self._previous_backoff = 0.0
 
     def reset(self) -> None:
-        self._previous_backoff = 0
+        self._previous_backoff = 0.0
 
     def compute(self, failures: int) -> float:
         max_backoff = max(self._base, self._previous_backoff * 3)

valkey/backoff.py

@@ -1,15 +1,13 @@
 from binascii import crc_hqx
 
-from valkey.typing import EncodedT
-
 # Valkey Cluster's key space is divided into 16384 slots.
 # For more information see: https://github.com/redis/redis/issues/2576
 VALKEY_CLUSTER_HASH_SLOTS = 16384
 
 __all__ = ["key_slot", "VALKEY_CLUSTER_HASH_SLOTS"]
 
 
-def key_slot(key: EncodedT, bucket: int = VALKEY_CLUSTER_HASH_SLOTS) -> int:
+def key_slot(key: bytes, bucket: int = VALKEY_CLUSTER_HASH_SLOTS) -> int:
     """Calculate key slot for a given key.
     See Keys distribution model in https://redis.io/topics/cluster-spec
     :param key - bytes

valkey/crc.py

@@ -17,9 +17,13 @@ class Lock:
     multiple clients play nicely together.
     """
 
-    lua_release = None
-    lua_extend = None
-    lua_reacquire = None
+    # This function is used to flag that the class `lua_*` functions are not set yet.
+    # Using a function, rather than `None`, prevents type annotation issues.
+    def __undefined(*_, **__) -> None: ...
+
+    lua_release = __undefined
+    lua_extend = __undefined
+    lua_reacquire = __undefined
 
     # KEYS[1] - lock name
     # ARGV[1] - token
@@ -147,11 +151,11 @@ def __init__(
     def register_scripts(self) -> None:
         cls = self.__class__
         client = self.valkey
-        if cls.lua_release is None:
+        if cls.lua_release is cls.__undefined:
             cls.lua_release = client.register_script(cls.LUA_RELEASE_SCRIPT)
-        if cls.lua_extend is None:
+        if cls.lua_extend is cls.__undefined:
             cls.lua_extend = client.register_script(cls.LUA_EXTEND_SCRIPT)
-        if cls.lua_reacquire is None:
+        if cls.lua_reacquire is cls.__undefined:
             cls.lua_reacquire = client.register_script(cls.LUA_REACQUIRE_SCRIPT)
 
     def __enter__(self) -> "Lock":
@@ -175,7 +179,7 @@ def acquire(
         sleep: Optional[Number] = None,
         blocking: Optional[bool] = None,
         blocking_timeout: Optional[Number] = None,
-        token: Optional[str] = None,
+        token: str| bytes | None = None,
     ):
         """
         Use Valkey to hold a shared, distributed lock named ``name``.
@@ -195,10 +199,10 @@ def acquire(
         if sleep is None:
             sleep = self.sleep
         if token is None:
-            token = uuid.uuid1().hex.encode()
+            encoded_token = uuid.uuid1().hex.encode()
         else:
             encoder = self.valkey.get_encoder()
-            token = encoder.encode(token)
+            encoded_token = encoder.encode(token)
         if blocking is None:
             blocking = self.blocking
         if blocking_timeout is None:
@@ -207,8 +211,8 @@ def acquire(
         if blocking_timeout is not None:
             stop_trying_at = mod_time.monotonic() + blocking_timeout
         while True:
-            if self.do_acquire(token):
-                self.local.token = token
+            if self.do_acquire(encoded_token):
+                self.local.token = encoded_token
                 return True
             if not blocking:
                 return False
@@ -217,7 +221,7 @@ def acquire(
                 return False
             mod_time.sleep(sleep)
 
-    def do_acquire(self, token: str) -> bool:
+    def do_acquire(self, token: bytes) -> bool:
         if self.timeout:
             # convert to milliseconds
             timeout = int(self.timeout * 1000)
@@ -312,6 +316,10 @@ def reacquire(self) -> bool:
         return self.do_reacquire()
 
     def do_reacquire(self) -> bool:
+        # `do_reacquire()` will only be called if `self.timeout` is not `None`.
+        # However, this `assert` is needed so that mypy understands the type.
+        assert self.timeout is not None
+
         timeout = int(self.timeout * 1000)
         if not bool(
             self.lua_reacquire(

valkey/lock.py

@@ -189,8 +189,7 @@ def owns_connection(self, connection):
         check = not self.is_master or (
             self.is_master and self.master_address == (connection.host, connection.port)
         )
-        parent = super()
-        return check and parent.owns_connection(connection)
+        return check and super().owns_connection(connection)
 
     def get_master_address(self):
         return self.proxy.get_master_address()

valkey/sentinel.py

@@ -63,7 +63,7 @@ def dict_merge(*dicts: Mapping[str, Any]) -> Dict[str, Any]:
     *dicts : `dict`
         dictionaries to merge
     """
-    merged = {}
+    merged: dict[str, Any] = {}
 
     for d in dicts:
         merged.update(d)