typings/faiss/__init__.pyi

"""Minimal type stubs for faiss used in this project.

This file only exposes a small subset of the faiss API that the
project uses, including the runtime-monkeypatched signatures such as
`Index.add_with_ids` so Pyright/Pylance stops reporting false positives.
"""

from typing import Any, overload

import numpy as np

class Index:
    d: int
    ntotal: int
    code_size: int
    nprobe: int

    def add(self, x: np.ndarray) -> None: ...
    def add_with_ids(self, x: np.ndarray, ids: np.ndarray) -> None: ...
    def search(
        self,
        x: np.ndarray,
        k: int,
        *,
        params: Any = ...,
        D: np.ndarray | None = ...,
        I: np.ndarray | None = ...,
    ) -> tuple[np.ndarray, np.ndarray]: ...
    def remove_ids(self, x: np.ndarray) -> int: ...
    @overload
    def reconstruct(self, key: int) -> np.ndarray: ...
    @overload
    def reconstruct(self, key: int, x: np.ndarray) -> None: ...
    def reconstruct(
        self, key: int, x: np.ndarray | None = ...
    ) -> np.ndarray | None: ...
    @overload
    def reconstruct_n(self, n0: int, ni: int) -> np.ndarray: ...
    @overload
    def reconstruct_n(self, n0: int, ni: int, x: np.ndarray) -> None: ...
    def reconstruct_n(
        self, n0: int = ..., ni: int = ..., x: np.ndarray | None = ...
    ) -> np.ndarray | None: ...
    def range_search(
        self, x: np.ndarray, thresh: float, *, params: Any = ...
    ) -> tuple[np.ndarray, np.ndarray, np.ndarray]: ...
    def add_sa_codes(self, codes: np.ndarray, ids: np.ndarray | None = ...) -> None: ...
    def sa_encode(self, x: np.ndarray) -> np.ndarray: ...
    def sa_decode(self, codes: np.ndarray) -> np.ndarray: ...

class IndexFlatL2(Index):
    def __init__(self, d: int) -> None: ...

class IndexIDMap(Index):
    index: Index

    def __init__(self, index: Index) -> None: ...

def read_index(path: str) -> Index: ...
def write_index(index: Index, path: str | None = ...) -> None: ...
def normalize_L2(x: np.ndarray) -> None: ...

# Additional concrete-ish classes exposed by some faiss builds (SWIG helpers
# expose `downcast_*` helpers to convert generic objects to these concrete
# types). We keep these minimal — only the names are important for typing.
class IndexBinary(Index):
    def __init__(self, d: int) -> None: ...

class InvertedLists:
    def __len__(self) -> int: ...

class AdditiveQuantizer:
    pass

class Quantizer:
    pass

class VectorTransform:
    pass

# SWIG-provided downcast helpers (present in some faiss Python builds).
def downcast_IndexBinary(obj: Any) -> IndexBinary: ...
def downcast_InvertedLists(obj: Any) -> InvertedLists: ...
def downcast_AdditiveQuantizer(obj: Any) -> AdditiveQuantizer: ...
def downcast_Quantizer(obj: Any) -> Quantizer: ...
def downcast_VectorTransform(obj: Any) -> VectorTransform: ...
def downcast_index(obj: Any) -> Index: ...

# version exposed by runtime
__version__: str