Skip to content

Add type annotations #611

@lbhm

Description

@lbhm

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: ...

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Type

    No type

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions