diff --git a/pyrightconfig.stricter.json b/pyrightconfig.stricter.json index f1ff078f3..e8dbd2729 100644 --- a/pyrightconfig.stricter.json +++ b/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", diff --git a/stubs/hnswlib/@tests/stubtest_allowlist.txt b/stubs/hnswlib/@tests/stubtest_allowlist.txt new file mode 100644 index 000000000..c8b95d1e5 --- /dev/null +++ b/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 diff --git a/stubs/hnswlib/METADATA.toml b/stubs/hnswlib/METADATA.toml new file mode 100644 index 000000000..67c290fb0 --- /dev/null +++ b/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" diff --git a/stubs/hnswlib/hnswlib.pyi b/stubs/hnswlib/hnswlib.pyi new file mode 100644 index 000000000..5f6bc2cdf --- /dev/null +++ b/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: ...