hnswlib icon indicating copy to clipboard operation
hnswlib copied to clipboard

Add type annotations

Open lbhm opened this issue 11 months ago • 1 comments

It would be great to add type annotations to the library so that hnswlib plays nice with static type checkers.

I already created type annotations for some functions myself (see below) and can offer to draft the other annotations as well. As far as I know, all it takes is a py.typed marker file and a hnswlib.pyi file in python_bindings/.

Initial type annotations:

from typing import Literal

import numpy as np
from numpy.typing import NDArray

class Index:
    def __init__(self, space: Literal["l2", "ip", "cosine"], dim: int) -> None: ...
    def init_index(
        self,
        max_elements: int,
        M: int = 16,  # noqa: N803
        ef_construction: int = 200,
        random_seed: int = 100,
        allow_replace_delete: bool = False,
    ) -> None: ...
    def add_items(
        self,
        data: NDArray[np.float32],
        ids: NDArray[np.uint64],
        num_threads: int = -1,
        replace_deleted: bool = False,
    ) -> None: ...
    def set_ef(self, ef: int) -> None: ...
    def knn_query(
        self,
        data: NDArray[np.float32],
        k: int = 1,
        num_threads: int = -1,
        filter: NDArray[np.uint64] | None = None,  # noqa: A002
    ) -> 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 save_index(self, path_to_index: str) -> None: ...

lbhm avatar Feb 12 '25 16:02 lbhm

For anyone interested in using type annotations: I contributed hnswlib annotations to typeshed so you can install them using pip install types-hnswlib.

lbhm avatar Feb 26 '25 06:02 lbhm