simplify vector type for Annoy (#5455)

Sequence was unnecessarily specific and prevents numpy type stubs from working with this package. Numpy ndarray isn't compatible with Sequence since it doesn't have __reversed__ method.
This commit is contained in:
Christopher Dignam
2021-05-14 20:29:31 -04:00
committed by GitHub
parent 15b0959cf8
commit 3aa4c67839

View File

@@ -1,5 +1,8 @@
from typing import Sequence, overload
from typing_extensions import Literal
from typing import Sized, overload
from typing_extensions import Literal, Protocol
class _Vector(Protocol, Sized):
def __getitem__(self, i: int) -> float: ...
class AnnoyIndex:
f: int
@@ -18,18 +21,18 @@ class AnnoyIndex:
) -> tuple[list[int], list[float]]: ...
@overload
def get_nns_by_vector(
self, vector: Sequence[float], n: int, search_k: int = ..., include_distances: Literal[False] = ...
self, vector: _Vector, n: int, search_k: int = ..., include_distances: Literal[False] = ...
) -> list[int]: ...
@overload
def get_nns_by_vector(
self, vector: Sequence[float], n: int, search_k: int, include_distances: Literal[True]
self, vector: _Vector, n: int, search_k: int, include_distances: Literal[True]
) -> tuple[list[int], list[float]]: ...
@overload
def get_nns_by_vector(
self, vector: Sequence[float], n: int, search_k: int = ..., *, include_distances: Literal[True]
self, vector: _Vector, n: int, search_k: int = ..., *, include_distances: Literal[True]
) -> tuple[list[int], list[float]]: ...
def get_item_vector(self, __i: int) -> list[float]: ...
def add_item(self, i: int, vector: Sequence[float]) -> None: ...
def add_item(self, i: int, vector: _Vector) -> None: ...
def on_disk_build(self, fn: str) -> Literal[True]: ...
def build(self, n_trees: int, n_jobs: int = ...) -> Literal[True]: ...
def unbuild(self) -> Literal[True]: ...