Add type annotations for hnswlib (#13529)

Author: lbhm

pyrightconfig.stricter.json

@@ -52,6 +52,7 @@
         "stubs/geopandas",
         "stubs/google-cloud-ndb",
         "stubs/hdbcli/hdbcli/dbapi.pyi",
+        "stubs/hnswlib",
         "stubs/html5lib",
         "stubs/httplib2",
         "stubs/humanfriendly",

stubs/hnswlib/@tests/stubtest_allowlist.txt

@@ -0,0 +1,3 @@
+# The metaclass of Index is pybind11_builtins.pybind11_type. Since pybind11 is not a
+# runtime dependency of hnswlib, we exclude Index from stubtest checks.
+hnswlib.Index

stubs/hnswlib/METADATA.toml

@@ -0,0 +1,4 @@
+version = "0.8.*"
+# Requires a version of numpy with a `py.typed` file
+requires = ["numpy>=1.20"]
+upstream_repository = "https://github.com/nmslib/hnswlib"

stubs/hnswlib/hnswlib.pyi

@@ -0,0 +1,62 @@
+from _typeshed import Incomplete
+from collections.abc import Callable
+from typing import Any, Literal, overload
+
+import numpy as np
+from numpy.typing import NDArray
+
+BFIndex: Incomplete
+
+class Index:
+    ef: int
+    num_threads: int
+
+    @overload
+    def __init__(self, params: dict[str, Any]) -> None: ...
+    @overload
+    def __init__(self, index: Index) -> None: ...
+    @overload
+    def __init__(self, space: Literal["l2", "ip", "cosine"], dim: int) -> None: ...
+    def add_items(self, data, ids: Incomplete | None = None, num_threads: int = -1, replace_deleted: bool = False) -> None: ...
+    def get_current_count(self) -> int: ...
+    def get_ids_list(self) -> list[int]: ...
+    @overload
+    def get_items(self, ids: Incomplete | None = ..., return_type: Literal["list"] = ...) -> list[float]: ...
+    @overload
+    def get_items(self, ids: Incomplete | None = ..., return_type: Literal["numpy"] = ...) -> NDArray[np.float32]: ...
+    @overload
+    def get_items(
+        self, ids: Incomplete | None = None, return_type: Literal["numpy", "list"] = "numpy"
+    ) -> NDArray[np.float32] | list[float]: ...
+    def get_max_elements(self) -> int: ...
+    def index_file_size(self) -> int: ...
+    def init_index(
+        self,
+        max_elements: int,
+        M: int = 16,
+        ef_construction: int = 200,
+        random_seed: int = 100,
+        allow_replace_delete: bool = False,
+    ) -> None: ...
+    def knn_query(
+        self, data, k: int = 1, num_threads: int = -1, filter: Callable[[int], bool] | None = None
+    ) -> tuple[NDArray[np.uint64], NDArray[np.float32]]: ...
+    def load_index(self, path_to_index: str, max_elements: int = 0, allow_replace_delete: bool = False) -> None: ...
+    def mark_deleted(self, label: int) -> None: ...
+    def resize_index(self, new_size: int) -> None: ...
+    def save_index(self, path_to_index: str) -> None: ...
+    def set_ef(self, ef: int) -> None: ...
+    def set_num_threads(self, num_threads: int) -> None: ...
+    def unmark_deleted(self, label: int) -> None: ...
+    @property
+    def M(self) -> int: ...
+    @property
+    def dim(self) -> int: ...
+    @property
+    def ef_construction(self) -> int: ...
+    @property
+    def element_count(self) -> int: ...
+    @property
+    def max_elements(self) -> int: ...
+    @property
+    def space(self) -> str: ...