chore: add pyright type checking and ci checks for it (#280)

* add config to pyproject toml

* add checks for pyright

* update

* update

* fix: add type ignore comments for Redis storage methods and update type hints in MinHash and LSH classes

* fix: update input validation in WeightedMinHashGenerator to check for sized iterables

* fix: add type hints and improve error handling in storage functions

* fix: add assertions to ensure keys and hashtables are initialized in AsyncMinHashLSH

* fix: update _merge method signature to return None and improve error handling in _hashed_byteswap

* fix: remove unnecessary assertion for upper bound in MinHashLSHEnsemble

* fix: simplify hashvalues initialization in WeightedMinHashGenerator

* fix: remove unnecessary assertions for keys and hashtables in AsyncMinHashLSH

* refactor: remove duplicate parsing of hashvalues

* refactor: update type hints for hashvalues and permutations in MinHash class

* apply suggestion

Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>

* revert seeds none change

---------

Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>

Author: bhimrazy

.github/workflows/checks.yml

@@ -10,6 +10,21 @@ jobs:
       - uses: actions/checkout@v6
       - uses: astral-sh/ruff-action@v3
 
+  pyright:
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v6
+      - uses: astral-sh/setup-uv@v7
+        with:
+          activate-environment: true
+          python-version: "3.10"
+          enable-cache: true
+      - name: Test Pyright
+        run: |
+          uvx pyright
+      - name: Minimize uv cache
+        run: uv cache prune --ci
+
   link-check:
     runs-on: ubuntu-latest
     steps:

datasketch/experimental/aio/storage.py

@@ -332,10 +332,10 @@ def initialized(self):
 
     class AsyncRedisListStorage(OrderedStorage, AsyncRedisStorage):
         async def keys(self):
-            return await self._redis.hkeys(self._name)
+            return await self._redis.hkeys(self._name)  # type: ignore
 
         async def redis_keys(self):
-            return await self._redis.hvals(self._name)
+            return await self._redis.hvals(self._name)  # type: ignore
 
         def status(self):
             status = self._parse_config(self.config["redis"])
@@ -374,8 +374,8 @@ async def remove_val(self, key, val, **kwargs):
                 await self._buffer.lrem(redis_key, val)
             else:
                 await self._redis.lrem(redis_key, val)
-                if not await self._redis.exists(redis_key):
-                    await self._redis.hdel(self._name, redis_key)
+                if not await self._redis.exists(redis_key):  # type: ignore
+                    await self._redis.hdel(self._name, redis_key)  # type: ignore
 
         async def insert(self, key, *vals, **kwargs):
             # Using buffer=True outside of an `insertion_session`
@@ -394,7 +394,7 @@ async def _insert(self, r, key, *values):
             await r.rpush(redis_key, *values)
 
         async def size(self):
-            return await self._redis.hlen(self._name)
+            return await self._redis.hlen(self._name)  # type: ignore
 
         async def itemcounts(self):
             pipe = self._redis.pipeline()
@@ -409,7 +409,7 @@ async def _get_len(r, k):
             return await r.llen(k)
 
         async def has_key(self, key):
-            return await self._redis.hexists(self._name, key)
+            return await self._redis.hexists(self._name, key)  # type: ignore
 
         async def empty_buffer(self):
             await self._buffer.execute()
@@ -429,8 +429,8 @@ async def remove_val(self, key, val, **kwargs):
                 await self._buffer.srem(redis_key, val)
             else:
                 await self._redis.srem(redis_key, val)
-                if not await self._redis.exists(redis_key):
-                    await self._redis.hdel(self._name, redis_key)
+                if not await self._redis.exists(redis_key):  # type: ignore
+                    await self._redis.hdel(self._name, redis_key)  # type: ignore
 
         async def _insert(self, r, key, *values):
             redis_key = self.redis_key(key)

datasketch/lsh.py

@@ -3,12 +3,18 @@
 import pickle
 import struct
 from collections.abc import Hashable
-from typing import Callable, Optional, Union
+from typing import Callable, List, Optional, Union
 
 from scipy.integrate import quad as integrate
 
 from datasketch.minhash import MinHash
-from datasketch.storage import _random_name, ordered_storage, unordered_storage
+from datasketch.storage import (
+    OrderedStorage,
+    UnorderedStorage,
+    _random_name,
+    ordered_storage,
+    unordered_storage,
+)
 from datasketch.weighted_minhash import WeightedMinHash
 
 
@@ -183,15 +189,15 @@ def __init__(
             self._H = self._byteswap
 
         basename = storage_config.get("basename", _random_name(11))
-        self.hashtables = [
+        self.hashtables: List[UnorderedStorage] = [
             unordered_storage(
                 storage_config,
                 name=b"".join([basename, b"_bucket_", struct.pack(">H", i)]),
             )
             for i in range(self.b)
         ]
         self.hashranges = [(i * self.r, (i + 1) * self.r) for i in range(self.b)]
-        self.keys = ordered_storage(storage_config, name=b"".join([basename, b"_keys"]))
+        self.keys: OrderedStorage = ordered_storage(storage_config, name=b"".join([basename, b"_keys"]))
 
     @property
     def buffer_size(self) -> int:
@@ -347,7 +353,7 @@ def __equivalent(self, other: MinHashLSH) -> bool:
         """
         return type(self) is type(other) and self.h == other.h and self.b == other.b and self.r == other.r
 
-    def _merge(self, other: MinHashLSH, check_overlap: bool = False, buffer: bool = False) -> MinHashLSH:
+    def _merge(self, other: MinHashLSH, check_overlap: bool = False, buffer: bool = False) -> None:
         if self.__equivalent(other):
             if check_overlap and set(self.keys).intersection(set(other.keys)):
                 raise ValueError("The keys are overlapping, duplicate key exists.")
@@ -524,6 +530,8 @@ def _byteswap(self, hs):
         return bytes(hs.byteswap().data)
 
     def _hashed_byteswap(self, hs):
+        if self.hashfunc is None:
+            raise RuntimeError("Hash function not configured.")
         return self.hashfunc(bytes(hs.byteswap().data))
 
     def _query_b(self, minhash, b):

datasketch/lsh_bloom.py

@@ -252,9 +252,9 @@ def __init__(
             raise ValueError("threshold must be in [0.0, 1.0]")
         if num_perm < 2:
             raise ValueError("Too few permutation functions")
-        if n <= 0:
+        if n is None or n <= 0:
             raise ValueError("n for LSHBloom must be >= 0")
-        if fp >= 1.0 or fp <= 0.0:
+        if fp is None or fp >= 1.0 or fp <= 0.0:
             raise ValueError("fp must be in (0.0, 1.0)")
         if save_dir is None:
             warnings.warn(

datasketch/lshensemble.py

@@ -221,7 +221,8 @@ def index(self, entries: Iterable[tuple[Hashable, MinHash, int]]) -> None:
         entries.sort(key=lambda e: e[2])
         curr_part = 0
         for key, minhash, size in entries:
-            if size > self.uppers[curr_part]:
+            u = self.uppers[curr_part]
+            if size > u:
                 curr_part += 1
             for r in self.indexes[curr_part]:
                 self.indexes[curr_part][r].insert(key, minhash)

datasketch/minhash.py

@@ -3,15 +3,18 @@
 import copy
 import warnings
 from collections.abc import Generator, Iterable
-from typing import Callable, Optional
+from typing import TYPE_CHECKING, Callable, Optional, Union
 
 try:
     from typing import Literal  # py3.8+; if older, you can fallback to typing_extensions
-except Exception:
+except ImportError:
     from typing_extensions import Literal
 
 import numpy as np
 
+if TYPE_CHECKING:
+    from numpy.typing import ArrayLike
+
 # GPU backend
 try:
     import cupy as cp
@@ -114,8 +117,8 @@ def __init__(
         gpu_mode: Literal["disable", "detect", "always"] = "disable",
         hashfunc: Callable = sha1_hash32,
         hashobj: Optional[object] = None,  # Deprecated.
-        hashvalues: Optional[Iterable] = None,
-        permutations: Optional[tuple[Iterable, Iterable]] = None,
+        hashvalues: Optional[ArrayLike] = None,
+        permutations: Optional[Union[tuple[ArrayLike, ArrayLike], ArrayLike]] = None,
     ) -> None:
         if hashvalues is not None:
             num_perm = len(hashvalues)
@@ -180,7 +183,7 @@ def _init_permutations(self, num_perm: int) -> np.ndarray:
             dtype=np.uint64,
         ).T
 
-    def _parse_hashvalues(self, hashvalues):
+    def _parse_hashvalues(self, hashvalues) -> np.ndarray:
         return np.array(hashvalues, dtype=np.uint64)
 
     def update(self, b) -> None:

datasketch/storage.py

@@ -26,7 +26,7 @@
     c_concurrent = None
 
 
-def ordered_storage(config, name=None):
+def ordered_storage(config, name=None) -> "OrderedStorage":
     """Return ordered storage system based on the specified config.
 
     The canonical example of such a storage container is
@@ -62,10 +62,10 @@ def ordered_storage(config, name=None):
         return RedisListStorage(config, name=name)
     if tp == "cassandra":
         return CassandraListStorage(config, name=name)
-    return None
+    raise ValueError(f"Unknown storage type: {tp}")
 
 
-def unordered_storage(config, name=None):
+def unordered_storage(config, name=None) -> "UnorderedStorage":
     """Return an unordered storage system based on the specified config.
 
     The canonical example of such a storage container is
@@ -100,7 +100,7 @@ def unordered_storage(config, name=None):
         return RedisSetStorage(config, name=name)
     if tp == "cassandra":
         return CassandraSetStorage(config, name=name)
-    return None
+    raise ValueError(f"Unknown storage type: {tp}")
 
 
 class Storage(ABC):
@@ -144,7 +144,7 @@ def insert(self, key, *vals, **kwargs):
         pass
 
     @abstractmethod
-    def remove(self, *keys):
+    def remove(self, *keys, **kwargs):
         """Remove `keys` from storage."""
         pass
 
@@ -154,12 +154,12 @@ def remove_val(self, key, val):
         pass
 
     @abstractmethod
-    def size(self):
+    def size(self) -> int:
         """Return size of storage with respect to number of keys."""
         pass
 
     @abstractmethod
-    def itemcounts(self, **kwargs):
+    def itemcounts(self, **kwargs) -> dict:
         """Returns the number of items stored under each key."""
         pass
 
@@ -168,6 +168,14 @@ def has_key(self, key):
         """Determines whether the key is in the storage or not."""
         pass
 
+    @property
+    def buffer_size(self) -> int:
+        return getattr(self, "_buffer_size", 50000)
+
+    @buffer_size.setter
+    def buffer_size(self, value: int):
+        self._buffer_size = value
+
     def status(self):
         return {"keyspace_size": len(self)}
 

datasketch/weighted_minhash.py

@@ -133,14 +133,15 @@ def minhash(self, v: np.ndarray) -> WeightedMinHash:
             WeightedMinHash: The weighted MinHash.
 
         """
-        if not isinstance(v, collections.abc.Iterable):
-            raise TypeError("Input vector must be an iterable")
+        if not isinstance(v, collections.abc.Sized):
+            raise TypeError("Input vector must be sized")
         if not len(v) == self.dim:
             raise ValueError("Input dimension mismatch, expecting %d" % self.dim)
         if not isinstance(v, np.ndarray):
             v = np.array(v, dtype=np.float32)
         elif v.dtype != np.float32:
             v = v.astype(np.float32)
+        v: np.ndarray = v
         hashvalues = np.zeros((self.sample_size, 2), dtype=int)
         vzeros = v == 0
         if vzeros.all():
@@ -226,9 +227,8 @@ def minhash_many(self, X: Union[sp.sparse.spmatrix, np.ndarray]) -> list[Union[W
                 doc_argmin = np.argmin(doc_ln_a, axis=1)
                 doc_k = doc_cidx[doc_argmin]
 
-                all_hashvalues[it_doc] = np.zeros((self.sample_size, 2), dtype=int)
-
-                hashvalues = all_hashvalues[it_doc]
+                hashvalues = np.zeros((self.sample_size, 2), dtype=int)
+                all_hashvalues[it_doc] = hashvalues
                 hashvalues[:, 0], hashvalues[:, 1] = (
                     doc_k,
                     t[np.arange(self.sample_size), doc_begin + doc_argmin],

pyproject.toml

@@ -164,6 +164,34 @@ addopts = ["--strict-markers", "--color=yes", "--cov-report=xml"]
 testpaths = ["test"]
 asyncio_mode = "auto"
 
+[tool.pyright]
+include = ["datasketch"]
+exclude = [
+  "benchmark",
+  "docs",
+  "examples",
+  "test",
+  "travis",
+  "**/.venv/**",
+  "**/__pycache__",
+]
+pythonVersion = "3.9"
+typeCheckingMode = "basic" # todo: change to "strict" in future
+
+reportMissingImports = "none"
+reportUnusedVariable = "warning"
+reportAttributeAccessIssue = "none"
+reportOptionalMemberAccess = "none"
+reportGeneralTypeIssues = "none"
+reportArgumentType = "none"
+reportOptionalIterable = "none"
+reportReturnType = "none"
+reportRedeclaration = "none"
+reportOperatorIssue = "none"
+reportAssignmentType = "none"
+reportOptionalSubscript = "none"
+reportCallIssue = "none"
+
 [tool.coverage.run]
 source = ["datasketch"]
 omit = ["*/experimental/*", "*/tests/*", "*/test/*"]