From 439c0b62cd42921c81250bd75319059c4a3410c9 Mon Sep 17 00:00:00 2001 From: Roman Solomatin <36135455+Samoed@users.noreply.github.com> Date: Sun, 29 Mar 2026 13:26:24 +0300 Subject: [PATCH 1/6] update pyproject to include ruff rules directly and groups --- CONTRIBUTING.md | 6 +++--- pyproject.toml | 26 +++++++++++++++++++++----- 2 files changed, 24 insertions(+), 8 deletions(-) diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index c4cfc937f..c69315086 100644 --- a/CONTRIBUTING.md +++ b/CONTRIBUTING.md @@ -214,7 +214,7 @@ Following options are enabled: - The `-p no:warnings` option will suppress and allow warnings. ```sh -uv pip install pytest pytest-repeat numpy # for repeated fuzzy tests +uv pip install -e . --group tests # for repeated fuzzy tests python -m pytest # if you trust the default settings python -m pytest python/scripts/ -s -x -p no:warnings # to overwrite the default settings ``` @@ -222,8 +222,8 @@ python -m pytest python/scripts/ -s -x -p no:warnings # to overwrite the default Linting: ```sh -pip install ruff -ruff --format=github --select=E9,F63,F7,F82 --target-version=py310 python +uv pip install -e . --group lint +ruff --format=github python ``` Before merging your changes you may want to test your changes against the entire matrix of Python versions USearch supports. diff --git a/pyproject.toml b/pyproject.toml index bdbf2ad4f..78afe30d1 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -86,8 +86,24 @@ before-build = [ "git submodule update --init --recursive", ] -# Configuration options for the Black formatter: -# https://black.readthedocs.io/en/latest/usage_and_configuration/the_basics.html#where-black-looks-for-the-file -[tool.black] -line-length = 120 # Set line length to the same value as in `.clang-format` for modern wide screens -target-version = ['py36', 'py314'] + +[dependency-groups] +lint = [ + "ruff>=0.15.8", +] +tests = [ + "numpy>=2.4.3", + "pytest>=9.0.2", + "pytest-repeat>=0.9.4", +] + +[tool.ruff] +target-version = "py310" + +[tool.ruff.lint] +select = [ + "E9", + "F63", + "F7", + "F82", +] From cfe901a609612b26b5f0a3460a9a73cf03251bf4 Mon Sep 17 00:00:00 2001 From: Roman Solomatin <36135455+Samoed@users.noreply.github.com> Date: Sun, 29 Mar 2026 15:33:56 +0300 Subject: [PATCH 2/6] add typing --- pyproject.toml | 23 ++- python/usearch/__init__.py | 48 +++-- python/usearch/client.py | 68 ++++--- python/usearch/eval.py | 226 ++++++++++++++-------- python/usearch/index.py | 385 +++++++++++++++++++++---------------- python/usearch/io.py | 31 ++- python/usearch/numba.py | 20 +- python/usearch/server.py | 48 +++-- setup.py | 6 + 9 files changed, 529 insertions(+), 326 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 78afe30d1..d2c0f9a71 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -92,10 +92,16 @@ lint = [ "ruff>=0.15.8", ] tests = [ - "numpy>=2.4.3", + "numpy>=2", "pytest>=9.0.2", "pytest-repeat>=0.9.4", ] +typing = [ + "mypy>=1.19.1", + "types-networkx>=3.6.1.20260321", + "types-setuptools>=82.0.0.20260210", + "types-tqdm>=4.67.3.20260303", +] [tool.ruff] target-version = "py310" @@ -106,4 +112,19 @@ select = [ "F63", "F7", "F82", + "TC", + "T20", + "UP", + "I", +] +future-annotations = true # For TC rules + +[tool.mypy] +python_version = "3.10" + +[[tool.mypy.overrides]] +# these modules not typed and don't have stubs +module = [ + "matplotlib", ] +ignore_missing_imports = true diff --git a/python/usearch/__init__.py b/python/usearch/__init__.py index e5f559e89..f220d3879 100644 --- a/python/usearch/__init__.py +++ b/python/usearch/__init__.py @@ -1,10 +1,9 @@ -import os -import sys import ctypes +import os import platform -import warnings +import sys import urllib.request -from typing import Optional, Tuple +import warnings from urllib.error import HTTPError #! Load SimSIMD before the USearch compiled module @@ -30,26 +29,26 @@ pass # If the user doesn't want SimSIMD, we assume they know what they're doing -from usearch.compiled import ( - VERSION_MAJOR, - VERSION_MINOR, - VERSION_PATCH, +from usearch.compiled import ( # type: ignore[import-not-found] # Default values: DEFAULT_CONNECTIVITY, DEFAULT_EXPANSION_ADD, DEFAULT_EXPANSION_SEARCH, + USES_FP16LIB, # Dependencies: USES_OPENMP, - USES_FP16LIB, USES_SIMSIMD, USES_SIMSIMD_DYNAMIC_DISPATCH, + VERSION_MAJOR, + VERSION_MINOR, + VERSION_PATCH, ) __version__ = f"{VERSION_MAJOR}.{VERSION_MINOR}.{VERSION_PATCH}" class BinaryManager: - def __init__(self, version: Optional[str] = None): + def __init__(self, version: str | None = None): if version is None: version = __version__ self.version = version or __version__ @@ -72,7 +71,7 @@ def determine_download_url(version: str, filename: str) -> str: url = f"{base_url}/v{version}/{filename}" return url - def get_binary_name(self) -> Tuple[str, str]: + def get_binary_name(self) -> tuple[str, str]: version = self.version os_map = {"Linux": "linux", "Windows": "windows", "Darwin": "macos"} arch_map = { @@ -85,12 +84,14 @@ def get_binary_name(self) -> Tuple[str, str]: os_part = os_map.get(platform.system(), "") arch = platform.machine() arch_part = arch_map.get(arch, "") - extension = {"Linux": "so", "Windows": "dll", "Darwin": "dylib"}.get(platform.system(), "") + extension = {"Linux": "so", "Windows": "dll", "Darwin": "dylib"}.get( + platform.system(), "" + ) source_filename = f"usearch_sqlite_{os_part}_{arch_part}_{version}.{extension}" target_filename = f"usearch_sqlite.{extension}" return source_filename, target_filename - def sqlite_found_or_downloaded(self) -> Optional[str]: + def sqlite_found_or_downloaded(self) -> str | None: """ Attempts to locate the pre-installed `usearch_sqlite` binary. If not found, downloads it from GitHub. @@ -104,7 +105,6 @@ def sqlite_found_or_downloaded(self) -> Optional[str]: # Check local development directories first for local_dir in local_dirs: - local_path = os.path.join(local_dir, target_filename) if os.path.exists(local_path): path_wout_extension, _, _ = local_path.rpartition(".") @@ -120,9 +120,10 @@ def sqlite_found_or_downloaded(self) -> Optional[str]: download_dir = self.determine_download_dir() local_path = os.path.join(download_dir, target_filename) if not os.path.exists(local_path): - # If not found locally, warn the user and download from GitHub - warnings.warn("Will download `usearch_sqlite` binary from GitHub.", UserWarning) + warnings.warn( + "Will download `usearch_sqlite` binary from GitHub.", UserWarning + ) try: source_url = self.determine_download_url(self.version, source_filename) os.makedirs(download_dir, exist_ok=True) @@ -130,9 +131,14 @@ def sqlite_found_or_downloaded(self) -> Optional[str]: except HTTPError as e: # If the download fails due to HTTPError (e.g., 404 Not Found), like a missing lib version if e.code == 404: - warnings.warn(f"Download failed: {e.url} could not be found.", UserWarning) + warnings.warn( + f"Download failed: {e.url} could not be found.", UserWarning + ) else: - warnings.warn(f"Download failed with HTTP error: {e.code} {e.reason}", UserWarning) + warnings.warn( + f"Download failed with HTTP error: {e.code} {e.reason}", + UserWarning, + ) return None # Handle the case where binary_path does not exist after supposed successful download @@ -140,11 +146,13 @@ def sqlite_found_or_downloaded(self) -> Optional[str]: path_wout_extension, _, _ = local_path.rpartition(".") return path_wout_extension else: - warnings.warn("Failed to download `usearch_sqlite` binary from GitHub.", UserWarning) + warnings.warn( + "Failed to download `usearch_sqlite` binary from GitHub.", UserWarning + ) return None -def sqlite_path(version: str = None) -> str: +def sqlite_path(version: str | None = None) -> str: manager = BinaryManager(version=version) result = manager.sqlite_found_or_downloaded() if result is None: diff --git a/python/usearch/client.py b/python/usearch/client.py index b8946b7ed..2c52bca3c 100644 --- a/python/usearch/client.py +++ b/python/usearch/client.py @@ -1,12 +1,17 @@ -from typing import Union, Optional, List +from __future__ import annotations + +from typing import TYPE_CHECKING, Any import numpy as np -from ucall.client import Client +from ucall.client import Client # type: ignore[import-untyped] + +if TYPE_CHECKING: + from numpy.typing import NDArray -from usearch.index import Matches +from usearch.index import BatchMatches, Matches -def _vector_to_ascii(vector: np.ndarray) -> Optional[str]: +def _vector_to_ascii(vector: NDArray[Any]) -> str | None: if vector.dtype != np.int8 and vector.dtype != np.uint8 and vector.dtype != np.byte: return None if not np.all((vector >= 0) | (vector <= 100)): @@ -21,10 +26,12 @@ def _vector_to_ascii(vector: np.ndarray) -> Optional[str]: class IndexClient: - def __init__(self, uri: str = "127.0.0.1", port: int = 8545, use_http: bool = True) -> None: + def __init__( + self, uri: str = "127.0.0.1", port: int = 8545, use_http: bool = True + ) -> None: self.client = Client(uri=uri, port=port, use_http=use_http) - def add_one(self, key: int, vector: np.ndarray): + def add_one(self, key: int, vector: NDArray[Any]): assert isinstance(key, int) assert isinstance(vector, np.ndarray) vector = vector.flatten() @@ -34,57 +41,57 @@ def add_one(self, key: int, vector: np.ndarray): else: self.client.add_one(key=key, vectors=vector) - def add_many(self, keys: np.ndarray, vectors: np.ndarray): + def add_many(self, keys: NDArray[Any], vectors: NDArray[Any]): assert isinstance(keys, int) assert isinstance(vectors, np.ndarray) assert keys.ndim == 1 and vectors.ndim == 2 assert keys.shape[0] == vectors.shape[0] self.client.add_many(keys=keys, vectors=vectors) - def add(self, keys: Union[np.ndarray, int], vectors: np.ndarray): + def add(self, keys: NDArray[Any] | int, vectors: NDArray[Any]): if isinstance(keys, int) or len(keys) == 1: - return self.add_one(keys, vectors) + return self.add_one( + int(keys) if isinstance(keys, np.ndarray) else keys, vectors + ) else: return self.add_many(keys, vectors) - def search_one(self, vector: np.ndarray, count: int) -> Matches: - matches: List[dict] = [] + def search_one(self, vector: NDArray[Any], count: int) -> Matches: vector = vector.flatten() ascii_vector = _vector_to_ascii(vector) if ascii_vector: - matches = self.client.search_ascii(string=ascii_vector, count=count) + raw = self.client.search_ascii(string=ascii_vector, count=count) else: - matches = self.client.search_one(vector=vector, count=count) + raw = self.client.search_one(vector=vector, count=count) - print(matches.data) - matches = matches.json + matches: list[dict] = raw.json - keys = np.array((1, count), dtype=np.uint32) - distances = np.array((1, count), dtype=np.float32) - counts = np.array((1), dtype=np.uint32) + keys = np.array(count, dtype=np.uint32) + distances = np.array(count, dtype=np.float32) for col, result in enumerate(matches): - keys[0, col] = result["key"] - distances[0, col] = result["distance"] - counts[0] = len(matches) + keys[col] = result["key"] + distances[col] = result["distance"] - return keys, distances, counts + return Matches(keys=keys[: len(matches)], distances=distances[: len(matches)]) - def search_many(self, vectors: np.ndarray, count: int) -> Matches: + def search_many(self, vectors: NDArray[Any], count: int) -> BatchMatches: batch_size: int = vectors.shape[0] - list_of_matches: List[List[dict]] = self.client.search_many(vectors=vectors, count=count) + list_of_matches: list[list[dict]] = self.client.search_many( + vectors=vectors, count=count + ) - keys = np.array((batch_size, count), dtype=np.uint32) - distances = np.array((batch_size, count), dtype=np.float32) - counts = np.array((batch_size), dtype=np.uint32) + keys = np.zeros((batch_size, count), dtype=np.uint32) + distances = np.zeros((batch_size, count), dtype=np.float32) + counts = np.zeros(batch_size, dtype=np.uint32) for row, matches in enumerate(list_of_matches): for col, result in enumerate(matches): keys[row, col] = result["key"] distances[row, col] = result["distance"] - counts[row] = len(results) + counts[row] = len(matches) - return keys, distances, counts + return BatchMatches(keys=keys, distances=distances, counts=counts) - def search(self, vectors: np.ndarray, count: int) -> Matches: + def search(self, vectors: NDArray[Any], count: int) -> Matches | BatchMatches: if vectors.ndim == 1 or (vectors.ndim == 2 and vectors.shape[0] == 1): return self.search_one(vectors, count) else: @@ -117,4 +124,3 @@ def save(self, path: str): index = IndexClient() index.add(42, np.array([0.4] * 256, dtype=np.float32)) results = index.search(np.array([0.4] * 256, dtype=np.float32), 10) - print(results) diff --git a/python/usearch/eval.py b/python/usearch/eval.py index c4e95122c..f943b9523 100644 --- a/python/usearch/eval.py +++ b/python/usearch/eval.py @@ -1,33 +1,39 @@ from __future__ import annotations -from time import time_ns -from typing import Tuple, Any, Callable, Union, Optional, List -from dataclasses import dataclass, asdict + from collections import defaultdict +from dataclasses import asdict, dataclass from math import ceil +from time import time_ns +from typing import TYPE_CHECKING, Any import numpy as np -from usearch.io import load_matrix from usearch.index import ( - Index, BatchMatches, - ScalarKind, + Index, + Key, MetricKind, MetricKindBitwise, - Key, - _normalize_metric, + ScalarKind, _normalize_dtype, + _normalize_metric, _to_numpy_dtype, ) +from usearch.io import load_matrix + +if TYPE_CHECKING: + from collections.abc import Callable + + from numpy.typing import NDArray def random_vectors( count: int, metric: MetricKind = MetricKind.IP, dtype: ScalarKind = ScalarKind.F32, - ndim: Optional[int] = None, - index: Optional[Index] = None, -) -> np.ndarray: + ndim: int | None = None, + index: Index | None = None, +) -> NDArray[Any]: """Produces a collection of random vectors normalized for the provided `metric` and matching wanted `dtype`, which can both be inferred from an existing `index`. """ @@ -38,18 +44,19 @@ def random_vectors( raise ValueError("Unsupported `index` type") ndim = index.ndim - dtype = index.numpy_dtype + dtype = index.dtype metric = index.metric else: - metric: MetricKind = _normalize_metric(metric) - dtype: ScalarKind = _normalize_dtype(dtype, ndim=ndim, metric=metric) + metric = _normalize_metric(metric) + dtype = _normalize_dtype(dtype, ndim=ndim or 0, metric=metric) + + assert ndim is not None, "ndim must be provided when index is not given" # Produce data if metric in MetricKindBitwise or dtype == ScalarKind.B1: - bit_vectors = np.random.randint(2, size=(count, ndim)) - bit_vectors = np.packbits(bit_vectors, axis=1) - return bit_vectors + bits = np.random.randint(2, size=(count, ndim)) + return np.packbits(bits, axis=1) else: x = np.random.rand(count, ndim) @@ -87,14 +94,16 @@ class SearchStats: @property def mean_efficiency(self) -> float: - return 1 - float(self.computed_distances) / (self.count_queries * self.index_size) + return 1 - float(self.computed_distances) / ( + self.count_queries * self.index_size + ) @property def mean_recall(self) -> float: return self.count_matches / self.count_queries -def self_recall(index: Index, sample: Union[float, int] = 1.0, **kwargs) -> SearchStats: +def self_recall(index: Index, sample: float | int = 1.0, **kwargs) -> SearchStats: """Simplest benchmark for a quality of search, which queries every existing member of the index, to make sure approximate search finds the point itself. @@ -107,7 +116,13 @@ def self_recall(index: Index, sample: Union[float, int] = 1.0, **kwargs) -> Sear :rtype: SearchStats """ if len(index) == 0: - return 0 + return SearchStats( + index_size=0, + count_queries=0, + count_matches=0, + visited_members=0, + computed_distances=0, + ) if "count" not in kwargs: kwargs["count"] = 1 @@ -128,7 +143,9 @@ def self_recall(index: Index, sample: Union[float, int] = 1.0, **kwargs) -> Sear matches = index.search(vectors, **kwargs) count_matches: int = ( - matches.count_matches(keys) if isinstance(matches, BatchMatches) else int(matches.keys[0] == keys[0]) + matches.count_matches(keys) + if isinstance(matches, BatchMatches) + else int(matches.keys[0] == keys[0]) ) return SearchStats( index_size=len(index), @@ -139,7 +156,7 @@ def self_recall(index: Index, sample: Union[float, int] = 1.0, **kwargs) -> Sear ) -def measure_seconds(f: Callable) -> Tuple[float, Any]: +def measure_seconds(f: Callable) -> tuple[float, Any]: """Simple function profiling decorator. :param f: Function to be profiled @@ -155,11 +172,11 @@ def measure_seconds(f: Callable) -> Tuple[float, Any]: return secs, result -def dcg(relevances: np.ndarray, k: Optional[int] = None) -> np.ndarray: +def dcg(relevances: NDArray[Any], k: int | None = None) -> float | np.floating[Any]: """Calculate DCG (Discounted Cumulative Gain) up to position k. :param relevances: List of true relevance scores (in the order as they are ranked) - :type relevances: list + :type relevances: np.ndarray :param k: Position up to which DCG is computed :type k: int :return: The DCG score at position k @@ -176,7 +193,7 @@ def dcg(relevances: np.ndarray, k: Optional[int] = None) -> np.ndarray: return np.sum(relevances / discounts) -def ndcg(relevances: np.ndarray, k: Optional[int] = None) -> np.ndarray: +def ndcg(relevances: NDArray[Any], k: int | None = None) -> float: """Calculate NDCG (Normalized Discounted Cumulative Gain) at position k. :param relevances: List of true relevance scores (in the order as they are ranked) @@ -186,14 +203,16 @@ def ndcg(relevances: np.ndarray, k: Optional[int] = None) -> np.ndarray: :return: The NDCG score at position k :rtype: float """ - best_dcg = dcg(sorted(relevances, reverse=True), k) + best_dcg = dcg(np.array(sorted(relevances, reverse=True)), k) if best_dcg == 0: return 0.0 - return dcg(relevances, k) / best_dcg + return float(dcg(relevances, k) / best_dcg) -def relevance(expected: np.ndarray, predicted: np.ndarray, k: Optional[int] = None) -> np.ndarray: +def relevance( + expected: NDArray[Any], predicted: NDArray[Any], k: int | None = None +) -> list[int]: """Calculate relevance scores. Binary relevance scores :param expected: ground-truth keys @@ -208,12 +227,13 @@ def relevance(expected: np.ndarray, predicted: np.ndarray, k: Optional[int] = No @dataclass class Dataset: - keys: np.ndarray - vectors: np.ndarray - queries: np.ndarray - neighbors: np.ndarray + keys: NDArray[Any] | None + vectors: NDArray[Any] | None + queries: NDArray[Any] | None + neighbors: NDArray[Any] | None def crop_neighbors(self, k: int): + assert self.neighbors is not None self.neighbors = self.neighbors[:, k] @property @@ -222,12 +242,12 @@ def ndim(self): @staticmethod def build( - vectors: Optional[str] = None, - queries: Optional[str] = None, - neighbors: Optional[str] = None, - count: Optional[int] = None, - ndim: Optional[int] = None, - k: Optional[int] = None, + vectors: str | None = None, + queries: str | None = None, + neighbors: str | None = None, + count: int | None = None, + ndim: int | None = None, + k: int | None = None, ): """Either loads an existing dataset from disk, or generates one on the fly. @@ -250,19 +270,33 @@ def build( if vectors is not None: assert ndim is None - d.vectors = load_matrix(vectors) + loaded_vectors = load_matrix(vectors) + assert loaded_vectors is not None, f"Failed to load vectors from {vectors}" + d.vectors = loaded_vectors ndim = d.vectors.shape[1] - count = min(d.vectors.shape[0], count) if count is not None else d.vectors.shape[0] + count = ( + min(d.vectors.shape[0], count) + if count is not None + else d.vectors.shape[0] + ) d.vectors = d.vectors[:count, :] d.keys = np.arange(count, dtype=Key) if queries is not None: - d.queries = load_matrix(queries) + loaded_queries = load_matrix(queries) + assert loaded_queries is not None, ( + f"Failed to load queries from {queries}" + ) + d.queries = loaded_queries else: d.queries = d.vectors if neighbors is not None: - d.neighbors = load_matrix(neighbors) + loaded_neighbors = load_matrix(neighbors) + assert loaded_neighbors is not None, ( + f"Failed to load neighbors from {neighbors}" + ) + d.neighbors = loaded_neighbors if k is not None: d.neighbors = d.neighbors[:, :k] else: @@ -284,12 +318,12 @@ def build( @dataclass class TaskResult: - add_operations: Optional[int] = None - add_per_second: Optional[float] = None + add_operations: int | None = None + add_per_second: float | None = None - search_operations: Optional[int] = None - search_per_second: Optional[float] = None - recall_at_one: Optional[float] = None + search_operations: int | None = None + search_per_second: float | None = None + recall_at_one: float | None = None def __repr__(self) -> str: parts = [] @@ -303,17 +337,21 @@ def __repr__(self) -> str: @property def add_seconds(self) -> float: + assert self.add_operations is not None and self.add_per_second is not None return self.add_operations / self.add_per_second @property def search_seconds(self) -> float: + assert self.search_operations is not None and self.search_per_second is not None return self.search_operations / self.search_per_second def __add__(self, other: TaskResult): result = TaskResult() if self.add_operations and other.add_operations: result.add_operations = self.add_operations + other.add_operations - result.add_per_second = result.add_operations / (self.add_seconds + other.add_seconds) + result.add_per_second = result.add_operations / ( + self.add_seconds + other.add_seconds + ) else: base = self if self.add_operations else other result.add_operations = base.add_operations @@ -322,9 +360,12 @@ def __add__(self, other: TaskResult): if self.search_operations and other.search_operations: result.search_operations = self.search_operations + other.search_operations result.recall_at_one = ( - self.recall_at_one * self.search_operations + other.recall_at_one * other.search_operations + (self.recall_at_one or 0.0) * self.search_operations + + (other.recall_at_one or 0.0) * other.search_operations ) / (self.search_operations + other.search_operations) - result.search_per_second = result.search_operations / (self.search_seconds + other.search_seconds) + result.search_per_second = result.search_operations / ( + self.search_seconds + other.search_seconds + ) else: base = self if self.search_operations else other result.search_operations = base.search_operations @@ -336,8 +377,8 @@ def __add__(self, other: TaskResult): @dataclass class AddTask: - keys: np.ndarray - vectors: np.ndarray + keys: NDArray[Any] + vectors: NDArray[Any] def __call__(self, index: Index) -> TaskResult: batch_size: int = self.vectors.shape[0] @@ -366,7 +407,7 @@ def inplace_shuffle(self): self.keys = self.keys[new_order] self.vectors = self.vectors[new_order, :] - def slices(self, batch_size: int) -> List[AddTask]: + def slices(self, batch_size: int) -> list[AddTask]: """Splits this dataset into smaller chunks.""" return [ @@ -377,10 +418,10 @@ def slices(self, batch_size: int) -> List[AddTask]: for start_row in range(0, self.count, batch_size) ] - def clusters(self, number_of_clusters: int) -> List[AddTask]: + def clusters(self, number_of_clusters: int) -> list[AddTask]: """Splits this dataset into smaller chunks.""" - from sklearn.cluster import KMeans + from sklearn.cluster import KMeans # type: ignore[import-untyped] clustering = KMeans( n_clusters=number_of_clusters, @@ -403,18 +444,20 @@ def clusters(self, number_of_clusters: int) -> List[AddTask]: @dataclass class SearchTask: - queries: np.ndarray - neighbors: np.ndarray + queries: NDArray[Any] + neighbors: NDArray[Any] def __call__(self, index: Index) -> TaskResult: - dt, results = measure_seconds(lambda: index.search(self.queries, self.neighbors.shape[1])) + dt, results = measure_seconds( + lambda: index.search(self.queries, self.neighbors.shape[1]) + ) return TaskResult( search_per_second=self.queries.shape[0] / dt, recall_at_one=results.mean_recall(self.neighbors[:, 0].flatten()), ) - def slices(self, batch_size: int) -> List[SearchTask]: + def slices(self, batch_size: int) -> list[SearchTask]: """Splits this dataset into smaller chunks.""" return [ @@ -428,13 +471,17 @@ def slices(self, batch_size: int) -> List[SearchTask]: @dataclass class Evaluation: - tasks: List[Union[AddTask, SearchTask]] + tasks: list[AddTask | SearchTask] count: int ndim: int @staticmethod - def for_dataset(dataset: Dataset, batch_size: int = 0, clusters: int = 1) -> Evaluation: - tasks = [] + def for_dataset( + dataset: Dataset, batch_size: int = 0, clusters: int = 1 + ) -> Evaluation: + tasks: list[AddTask | SearchTask] = [] + assert dataset.vectors is not None and dataset.keys is not None + assert dataset.queries is not None and dataset.neighbors is not None add = AddTask(vectors=dataset.vectors, keys=dataset.keys) search = SearchTask(queries=dataset.queries, neighbors=dataset.neighbors) @@ -443,7 +490,6 @@ def for_dataset(dataset: Dataset, batch_size: int = 0, clusters: int = 1) -> Eva tasks.extend(search.slices(batch_size)) elif clusters != 1: tasks.extend(add.clusters(clusters)) - print(tasks) tasks.append(search) else: tasks.append(add) @@ -476,18 +522,49 @@ def __call__(self, index: Index, post_clean: bool = True) -> dict: import argparse # Initialize the argument parser - parser = argparse.ArgumentParser(description="Evaluate vector search index for speed and accuracy.") + parser = argparse.ArgumentParser( + description="Evaluate vector search index for speed and accuracy." + ) # Define expected arguments - parser.add_argument("--vectors", type=str, required=False, help="Path to the file containing the vectors.") - parser.add_argument("--queries", type=str, required=False, help="Path to the file containing the query vectors.") - parser.add_argument("--neighbors", type=str, required=False, help="Path to the file with neighbor arrays.") - parser.add_argument("--dtype", type=str, required=False, help="Quantization type for internal storage.") + parser.add_argument( + "--vectors", + type=str, + required=False, + help="Path to the file containing the vectors.", + ) + parser.add_argument( + "--queries", + type=str, + required=False, + help="Path to the file containing the query vectors.", + ) + parser.add_argument( + "--neighbors", + type=str, + required=False, + help="Path to the file with neighbor arrays.", + ) + parser.add_argument( + "--dtype", + type=str, + required=False, + help="Quantization type for internal storage.", + ) parser.add_argument("--metric", type=str, required=False, help="Distance function.") parser.add_argument("--count", type=int, help="Number of vectors to use.") - parser.add_argument("--ndim", type=int, help="Number of dimensions for the vectors.") - parser.add_argument("--batch_size", type=int, default=0, help="Batch size for indexing and searching.") - parser.add_argument("--clusters", type=int, default=1, help="Number of clusters for indexing.") + parser.add_argument( + "--ndim", type=int, help="Number of dimensions for the vectors." + ) + parser.add_argument( + "--batch_size", + type=int, + default=0, + help="Batch size for indexing and searching.", + ) + parser.add_argument( + "--clusters", type=int, default=1, help="Number of clusters for indexing." + ) # Parse arguments from the command line args = parser.parse_args() @@ -502,11 +579,10 @@ def __call__(self, index: Index, post_clean: bool = True) -> dict: ) # Prepare the evaluation - evaluation = Evaluation.for_dataset(dataset, batch_size=args.batch_size, clusters=args.clusters) + evaluation = Evaluation.for_dataset( + dataset, batch_size=args.batch_size, clusters=args.clusters + ) index = Index(ndim=dataset.ndim, dtype=args.dtype, metric=args.metric) # Perform the evaluation results = evaluation(index) - - # Print the evaluation results - print("Evaluation results:", results) diff --git a/python/usearch/index.py b/python/usearch/index.py index 887bac223..505f370c2 100644 --- a/python/usearch/index.py +++ b/python/usearch/index.py @@ -1,55 +1,62 @@ from __future__ import annotations -from inspect import signature -from collections.abc import Sequence + +import math # The purpose of this file is to provide Pythonic wrapper on top # the native precompiled CPython module. It improves compatibility # Python tooling, linters, and static analyzers. It also embeds JIT # into the primary `Index` class, connecting USearch with Numba. import os -import sys -import math +from collections.abc import Callable, Iterable, Sequence from dataclasses import dataclass -from typing import ( - Any, - Optional, - Union, - NamedTuple, - List, - Iterable, - Tuple, - Dict, - Callable, -) +from inspect import signature +from typing import TYPE_CHECKING, Any, NamedTuple, TypeAlias, overload import numpy as np +from numpy.typing import NDArray from tqdm import tqdm +# Precompiled symbols that will be exposed +from usearch.compiled import ( # type: ignore[import-not-found] + DEFAULT_CONNECTIVITY, + DEFAULT_EXPANSION_ADD, + DEFAULT_EXPANSION_SEARCH, + USES_FP16LIB, + USES_OPENMP, + USES_SIMSIMD, + MetricKind, + MetricSignature, + ScalarKind, +) + # Precompiled symbols that won't be exposed directly: from usearch.compiled import ( Index as _CompiledIndex, +) +from usearch.compiled import ( Indexes as _CompiledIndexes, - IndexStats as _CompiledIndexStats, - index_dense_metadata_from_path as _index_dense_metadata_from_path, - index_dense_metadata_from_buffer as _index_dense_metadata_from_buffer, +) +from usearch.compiled import ( exact_search as _exact_search, +) +from usearch.compiled import ( hardware_acceleration as _hardware_acceleration, - kmeans as _kmeans, ) - -# Precompiled symbols that will be exposed from usearch.compiled import ( - MetricKind, - ScalarKind, - MetricSignature, - DEFAULT_CONNECTIVITY, - DEFAULT_EXPANSION_ADD, - DEFAULT_EXPANSION_SEARCH, - USES_OPENMP, - USES_SIMSIMD, - USES_FP16LIB, + index_dense_metadata_from_buffer as _index_dense_metadata_from_buffer, +) +from usearch.compiled import ( + index_dense_metadata_from_path as _index_dense_metadata_from_path, +) +from usearch.compiled import ( + kmeans as _kmeans, ) +if TYPE_CHECKING: + from usearch.compiled import ( + IndexStats as _CompiledIndexStats, # type: ignore[import-not-found] + ) + MetricKindBitwise = ( MetricKind.Hamming, MetricKind.Tanimoto, @@ -63,32 +70,30 @@ class CompiledMetric(NamedTuple): signature: MetricSignature -# Define TypeAlias for older Python versions -if sys.version_info >= (3, 10): - from typing import TypeAlias -else: - TypeAlias = object # Fallback for older Python versions - Key: TypeAlias = np.uint64 NoneType: TypeAlias = type(None) -KeyOrKeysLike = Union[Key, Iterable[Key], int, Iterable[int], np.ndarray, memoryview] +KeyOrKeysLike: TypeAlias = ( + Key | Iterable[Key] | int | Iterable[int] | NDArray[Any] | memoryview +) -VectorOrVectorsLike = Union[np.ndarray, Iterable[np.ndarray], memoryview] +VectorOrVectorsLike: TypeAlias = NDArray[Any] | Iterable[NDArray[Any]] | memoryview -DTypeLike = Union[str, ScalarKind] +DTypeLike: TypeAlias = str | ScalarKind -MetricLike = Union[str, MetricKind, CompiledMetric] +MetricLike: TypeAlias = str | MetricKind | CompiledMetric -BytesLike = Union[bytes, bytearray, memoryview] +BytesLike: TypeAlias = bytes | bytearray | memoryview -PathOrBuffer = Union[str, os.PathLike, BytesLike] +PathOrBuffer: TypeAlias = str | os.PathLike | BytesLike ProgressCallback = Callable[[int, int], bool] -def _match_signature(func: Callable[[Any], Any], arg_types: List[type], ret_type: type) -> bool: +def _match_signature( + func: Callable[..., Any], arg_types: list[type], ret_type: type +) -> bool: assert callable(func), "Not callable" sig = signature(func) param_types = [param.annotation for param in sig.parameters.values()] @@ -190,16 +195,18 @@ def _is_buffer(obj: Any) -> bool: def _search_in_compiled( compiled_callable: Callable, - vectors: np.ndarray, + vectors: VectorOrVectorsLike, *, - log: Union[str, bool], - progress: Optional[ProgressCallback], + log: str | bool, + progress: ProgressCallback | None, **kwargs, -) -> Union[Matches, BatchMatches]: +) -> Matches | BatchMatches: # assert isinstance(vectors, np.ndarray), "Expects a NumPy array" assert vectors.ndim == 1 or vectors.ndim == 2, "Expects a matrix or vector" - assert not progress or _match_signature(progress, [int, int], bool), "Invalid callback" + assert not progress or _match_signature(progress, [int, int], bool), ( + "Invalid callback" + ) if vectors.ndim == 1: vectors = vectors.reshape(1, len(vectors)) @@ -207,7 +214,7 @@ def _search_in_compiled( def distill_batch( batch_matches: BatchMatches, - ) -> Union[BatchMatches, Matches]: + ) -> BatchMatches | Matches: return batch_matches[0] if count_vectors == 1 else batch_matches progress_callback = progress @@ -249,12 +256,14 @@ def _add_to_compiled( vectors, copy: bool, threads: int, - log: Union[str, bool], - progress: Optional[ProgressCallback], -) -> Union[int, np.ndarray]: + log: str | bool, + progress: ProgressCallback | None, +) -> int | NDArray[Any]: # assert isinstance(vectors, np.ndarray), "Expects a NumPy array" - assert not progress or _match_signature(progress, [int, int], bool), "Invalid callback" + assert not progress or _match_signature(progress, [int, int], bool), ( + "Invalid callback" + ) assert vectors.ndim == 1 or vectors.ndim == 2, "Expects a matrix or vector" if vectors.ndim == 1: vectors = vectors.reshape(1, len(vectors)) @@ -315,8 +324,8 @@ def to_tuple(self) -> tuple: class Matches: """Search results for a single query.""" - keys: np.ndarray - distances: np.ndarray + keys: NDArray[Any] + distances: NDArray[Any] visited_members: int = 0 computed_distances: int = 0 @@ -333,9 +342,16 @@ def __getitem__(self, index: int) -> Match: else: raise IndexError(f"`index` must be an integer under {len(self)}") - def to_list(self) -> List[tuple]: + def __iter__(self): + for i in range(len(self)): + yield self[i] + + def to_list(self) -> list[tuple]: """Convert to list of (key, distance) tuples.""" - return [(int(key), float(distance)) for key, distance in zip(self.keys, self.distances)] + return [ + (int(key), float(distance)) + for key, distance in zip(self.keys, self.distances) + ] def __repr__(self) -> str: return f"usearch.Matches({len(self)})" @@ -350,15 +366,15 @@ class BatchMatches(Sequence): Attributes: keys: 2D array of shape (n_queries, k) containing match keys - distances: 2D array of shape (n_queries, k) containing distances + distances: 2D array of shape (n_queries, k) containing distances counts: 1D array of shape (n_queries,) with actual number of matches per query visited_members: Total graph nodes visited during search computed_distances: Total distance computations performed """ - keys: np.ndarray - distances: np.ndarray - counts: np.ndarray + keys: NDArray[Any] + distances: NDArray[Any] + counts: NDArray[Any] visited_members: int = 0 computed_distances: int = 0 @@ -366,7 +382,15 @@ class BatchMatches(Sequence): def __len__(self) -> int: return len(self.counts) - def __getitem__(self, index: int) -> Matches: + @overload + def __getitem__(self, index: int) -> Matches: ... + + @overload + def __getitem__(self, index: slice) -> list[Matches]: ... + + def __getitem__(self, index: int | slice) -> Matches | list[Matches]: + if isinstance(index, slice): + return [self.__getitem__(i) for i in range(*index.indices(len(self)))] if isinstance(index, int) and index < len(self): return Matches( keys=self.keys[index, : self.counts[index]], @@ -377,17 +401,17 @@ def __getitem__(self, index: int) -> Matches: else: raise IndexError(f"`index` must be an integer under {len(self)}") - def to_list(self) -> List[List[tuple]]: + def to_list(self) -> list[tuple]: """Flatten matches for all queries into a list of `(key, distance)` tuples.""" list_of_matches = [self.__getitem__(row) for row in range(self.__len__())] return [match.to_tuple() for matches in list_of_matches for match in matches] - def mean_recall(self, expected: np.ndarray, count: Optional[int] = None) -> float: + def mean_recall(self, expected: NDArray[Any], count: int | None = None) -> float: """Measures recall [0, 1] as of `Matches` that contain the corresponding `expected` entry anywhere among results.""" return self.count_matches(expected, count=count) / len(expected) - def count_matches(self, expected: np.ndarray, count: Optional[int] = None) -> int: + def count_matches(self, expected: NDArray[Any], count: int | None = None) -> int: """Measures recall [0, len(expected)] as of `Matches` that contain the corresponding `expected` entry anywhere among results. """ @@ -413,7 +437,7 @@ def __init__( self, index: Index, matches: BatchMatches, - queries: Optional[np.ndarray] = None, + queries: NDArray[Any] | None = None, ) -> None: if queries is None: queries = index._compiled.get_keys_in_slice() @@ -425,10 +449,10 @@ def __repr__(self) -> str: return f"usearch.Clustering(for {len(self.queries)} queries)" @property - def centroids_popularity(self) -> Tuple[np.ndarray, np.ndarray]: + def centroids_popularity(self) -> tuple[NDArray[Any], NDArray[Any]]: return np.unique(self.matches.keys, return_counts=True) - def members_of(self, centroid: Key) -> np.ndarray: + def members_of(self, centroid: Key) -> NDArray[Any]: return self.queries[self.matches.keys.flatten() == centroid] def subcluster(self, centroid: Key, **clustering_kwargs) -> Clustering: @@ -465,24 +489,24 @@ class IndexedKeys(Sequence): """View of all keys in the index.""" def __init__(self, index: Index) -> None: - self.index = index + self._idx = index def __len__(self) -> int: - return len(self.index) + return len(self._idx) - def __getitem__( + def __getitem__( # type: ignore[override] self, - offset_offsets_or_slice: Union[int, np.ndarray, slice], - ) -> Union[Key, np.ndarray]: + offset_offsets_or_slice: int | NDArray[Any] | slice, + ) -> Key | NDArray[Any]: if isinstance(offset_offsets_or_slice, slice): start, stop, step = offset_offsets_or_slice.indices(len(self)) if step != 1: raise ValueError("Slicing with a step is not supported") - return self.index._compiled.get_keys_in_slice(start, stop - start) + return self._idx._compiled.get_keys_in_slice(start, stop - start) elif isinstance(offset_offsets_or_slice, Iterable): offsets = np.array(offset_offsets_or_slice) - return self.index._compiled.get_keys_at_offsets(offsets) + return self._idx._compiled.get_keys_at_offsets(offsets) else: offset = int(offset_offsets_or_slice) @@ -490,18 +514,18 @@ def __getitem__( offset += len(self) if offset < 0 or offset >= len(self): raise IndexError("Index out of range") - return self.index._compiled.get_key_at_offset(offset) + return self._idx._compiled.get_key_at_offset(offset) - def __array__(self, dtype=None) -> np.ndarray: + def __array__(self, dtype=None) -> NDArray[Any]: if dtype is None: dtype = Key - return self.index._compiled.get_keys_in_slice().astype(dtype) + return self._idx._compiled.get_keys_in_slice().astype(dtype) class Index: """Fast approximate nearest neighbor search for dense vectors. - Supports various distance metrics (cosine, euclidean, inner product, etc.) + Supports various distance metrics (cosine, euclidean, inner product, etc.) and automatic precision optimization. Vector keys must be integers. All vectors must have the same dimensionality. @@ -516,12 +540,12 @@ def __init__( *, # All arguments must be named ndim: int = 0, metric: MetricLike = MetricKind.Cos, - dtype: Optional[DTypeLike] = None, - connectivity: Optional[int] = None, - expansion_add: Optional[int] = None, - expansion_search: Optional[int] = None, + dtype: DTypeLike | None = None, + connectivity: int | None = None, + expansion_add: int | None = None, + expansion_search: int | None = None, multi: bool = False, - path: Optional[os.PathLike] = None, + path: os.PathLike | None = None, view: bool = False, enable_key_lookups: bool = True, ) -> None: @@ -600,7 +624,9 @@ def __init__( self._metric_pointer = metric.pointer self._metric_signature = metric.signature else: - raise ValueError("The `metric` must be a `CompiledMetric` or a `MetricKind`") + raise ValueError( + "The `metric` must be a `CompiledMetric` or a `MetricKind`" + ) # Validate, that the right scalar type is defined dtype = _normalize_dtype(dtype, ndim, self._metric_kind) @@ -625,7 +651,7 @@ def __init__( self.load(path) @staticmethod - def metadata(path_or_buffer: PathOrBuffer) -> Optional[dict]: + def metadata(path_or_buffer: PathOrBuffer) -> dict | None: try: if _is_buffer(path_or_buffer): return _index_dense_metadata_from_buffer(path_or_buffer) @@ -638,7 +664,9 @@ def metadata(path_or_buffer: PathOrBuffer) -> Optional[dict]: raise e @staticmethod - def restore(path_or_buffer: PathOrBuffer, view: bool = False, **kwargs) -> Optional[Index]: + def restore( + path_or_buffer: PathOrBuffer, view: bool = False, **kwargs + ) -> Index | None: meta = Index.metadata(path_or_buffer) if not meta: return None @@ -661,14 +689,14 @@ def __len__(self) -> int: def add( self, - keys: KeyOrKeysLike, + keys: KeyOrKeysLike | None, vectors: VectorOrVectorsLike, *, copy: bool = True, threads: int = 0, - log: Union[str, bool] = False, - progress: Optional[ProgressCallback] = None, - ) -> Union[int, np.ndarray]: + log: str | bool = False, + progress: ProgressCallback | None = None, + ) -> int | NDArray[Any]: """Inserts one or move vectors into the index. For maximal performance the `keys` and `vectors` @@ -716,16 +744,16 @@ def search( *, threads: int = 0, exact: bool = False, - log: Union[str, bool] = False, - progress: Optional[ProgressCallback] = None, - ) -> Union[Matches, BatchMatches]: + log: str | bool = False, + progress: ProgressCallback | None = None, + ) -> Matches | BatchMatches: """Performs approximate nearest neighbors search for one or more queries. - + When searching with batch queries, returns BatchMatches that pre-allocates arrays for the requested `count` size. If fewer matches exist than requested (e.g., when count > index size), use individual query access via batch_matches[i] to get only valid results, or check batch_matches.counts to see actual result counts per query. - + :param vectors: Query vector or vectors. :type vectors: VectorOrVectorsLike :param count: Upper count on the number of matches to find @@ -758,16 +786,16 @@ def search( progress=progress, ) - def contains(self, keys: KeyOrKeysLike) -> Union[bool, np.ndarray]: + def contains(self, keys: KeyOrKeysLike) -> bool | NDArray[Any]: if isinstance(keys, Iterable): return self._compiled.contains_many(np.array(keys, dtype=Key)) else: return self._compiled.contains_one(int(keys)) - def __contains__(self, keys: KeyOrKeysLike) -> Union[bool, np.ndarray]: + def __contains__(self, keys: KeyOrKeysLike) -> bool | NDArray[Any]: return self.contains(keys) - def count(self, keys: KeyOrKeysLike) -> Union[int, np.ndarray]: + def count(self, keys: KeyOrKeysLike) -> int | NDArray[Any]: if isinstance(keys, Iterable): return self._compiled.count_many(np.array(keys, dtype=Key)) else: @@ -776,8 +804,8 @@ def count(self, keys: KeyOrKeysLike) -> Union[int, np.ndarray]: def get( self, keys: KeyOrKeysLike, - dtype: Optional[DTypeLike] = None, - ) -> Union[Optional[np.ndarray], Tuple[Optional[np.ndarray]]]: + dtype: DTypeLike | None = None, + ) -> Any: """Looks up one or more keys from the `Index`, retrieving corresponding vectors. Returns `None`, if one key is requested, and its not present. @@ -800,7 +828,9 @@ def get( dtype = _normalize_dtype(dtype) view_dtype = _to_numpy_dtype(dtype) if view_dtype is None: - raise NotImplementedError("The requested representation type is not supported by NumPy") + raise NotImplementedError( + "The requested representation type is not supported by NumPy" + ) def cast(result): if result is not None: @@ -809,17 +839,21 @@ def cast(result): is_one = not isinstance(keys, Iterable) if is_one: - keys = [keys] - if not isinstance(keys, np.ndarray): - keys = np.array(keys, dtype=Key) + actual_keys = np.array([keys], dtype=Key) + elif isinstance(keys, np.ndarray): + actual_keys = keys.astype(Key) else: - keys = keys.astype(Key) + actual_keys = np.array(list(keys), dtype=Key) # type: ignore[arg-type] - results = self._compiled.get_many(keys, dtype) - results = cast(results) if isinstance(results, np.ndarray) else [cast(result) for result in results] + results = self._compiled.get_many(actual_keys, dtype) + results = ( + cast(results) + if isinstance(results, np.ndarray) + else [cast(result) for result in results] + ) return results[0] if is_one else results - def __getitem__(self, keys: KeyOrKeysLike) -> Union[np.ndarray, Tuple[np.ndarray, np.ndarray]]: + def __getitem__(self, keys: KeyOrKeysLike) -> Any: """Looks up one or more keys from the `Index`, retrieving corresponding vectors. Returns `None`, if one key is requested, and its not present. @@ -840,7 +874,7 @@ def remove( *, compact: bool = False, threads: int = 0, - ) -> Union[int, np.ndarray]: + ) -> int | NDArray[Any]: """Removes one or move vectors from the index. When working with extremely large indexes, you may want to @@ -862,14 +896,14 @@ def remove( keys = np.array(keys, dtype=Key) return self._compiled.remove_many(keys, compact=compact, threads=threads) - def __delitem__(self, keys: KeyOrKeysLike) -> Union[int, np.ndarray]: + def __delitem__(self, keys: KeyOrKeysLike) -> int | NDArray[Any]: return self.remove(keys) def rename( self, from_: KeyOrKeysLike, to: KeyOrKeysLike, - ) -> Union[int, np.ndarray]: + ) -> int | NDArray[Any]: """Rename existing member vector or vectors. May be used in iterative clustering procedures, where one would iteratively @@ -893,7 +927,7 @@ def rename( return self._compiled.rename_many_to_one(from_, int(to)) else: - return self._compiled.rename_one_to_one(int(from_), int(to)) + return self._compiled.rename_one_to_one(int(from_), int(to)) # type: ignore[arg-type] @property def jit(self) -> bool: @@ -944,7 +978,7 @@ def serialized_length(self) -> int: return self._compiled.serialized_length @property - def metric_kind(self) -> Union[MetricKind, CompiledMetric]: + def metric_kind(self) -> MetricKind | CompiledMetric: """Returns the type of metric used for distance calculations. :return: The metric kind used in the index. @@ -953,7 +987,7 @@ def metric_kind(self) -> Union[MetricKind, CompiledMetric]: return self._metric_jit.kind if self._metric_jit else self._metric_kind @property - def metric(self) -> Union[MetricKind, CompiledMetric]: + def metric(self) -> MetricKind | CompiledMetric: """Returns the metric object used for distance calculations. :return: The metric used in the index. @@ -979,7 +1013,9 @@ def metric(self, metric: MetricLike): metric_pointer = metric.pointer metric_signature = metric.signature else: - raise ValueError("The `metric` must be a `CompiledMetric` or a `MetricKind`") + raise ValueError( + "The `metric` must be a `CompiledMetric` or a `MetricKind`" + ) return self._compiled.change_metric( metric_kind=metric_kind, @@ -1038,6 +1074,15 @@ def expansion_add(self) -> int: """ return self._compiled.expansion_add + @expansion_add.setter + def expansion_add(self, v: int): + """Sets the expansion parameter used during addition. + + :param v: The new expansion parameter for additions. + :type v: int + """ + self._compiled.expansion_add = v + @property def expansion_search(self) -> int: """Returns the expansion parameter used during searches. @@ -1049,15 +1094,6 @@ def expansion_search(self) -> int: """ return self._compiled.expansion_search - @expansion_add.setter - def expansion_add(self, v: int): - """Sets the expansion parameter used during addition. - - :param v: The new expansion parameter for additions. - :type v: int - """ - self._compiled.expansion_add = v - @expansion_search.setter def expansion_search(self, v: int): """Sets the expansion parameter used during searches. @@ -1069,9 +1105,9 @@ def expansion_search(self, v: int): def save( self, - path_or_buffer: Union[str, os.PathLike, NoneType] = None, - progress: Optional[ProgressCallback] = None, - ) -> Optional[bytes]: + path_or_buffer: str | os.PathLike | NoneType = None, + progress: ProgressCallback | None = None, + ) -> bytes | None: """Saves the index to a file or buffer. If `path_or_buffer` is not provided, it defaults to the path stored in `self.path`. @@ -1083,18 +1119,21 @@ def save( :return: The index data as bytes if saving to a buffer, otherwise None. :rtype: Optional[bytes] """ - assert not progress or _match_signature(progress, [int, int], bool), "Invalid callback signature" + assert not progress or _match_signature(progress, [int, int], bool), ( + "Invalid callback signature" + ) path_or_buffer = path_or_buffer if path_or_buffer is not None else self.path if path_or_buffer is None: return self._compiled.save_index_to_buffer(progress) else: self._compiled.save_index_to_path(os.fspath(path_or_buffer), progress) + return None def load( self, - path_or_buffer: Union[PathOrBuffer, NoneType] = None, - progress: Optional[ProgressCallback] = None, + path_or_buffer: PathOrBuffer | NoneType = None, + progress: ProgressCallback | None = None, ): """Loads the index from a file or buffer. @@ -1107,7 +1146,9 @@ def load( :raises Exception: If no source is defined. :raises RuntimeError: If the file does not exist. """ - assert not progress or _match_signature(progress, [int, int], bool), "Invalid callback signature" + assert not progress or _match_signature(progress, [int, int], bool), ( + "Invalid callback signature" + ) path_or_buffer = path_or_buffer if path_or_buffer is not None else self.path if path_or_buffer is None: @@ -1119,12 +1160,12 @@ def load( if os.path.exists(path_or_buffer): self._compiled.load_index_from_path(path_or_buffer, progress) else: - raise FileNotFoundError(f"File not found: {path_or_buffer}") + raise FileNotFoundError(f"File not found: {path_or_buffer!r}") def view( self, - path_or_buffer: Union[PathOrBuffer, NoneType] = None, - progress: Optional[ProgressCallback] = None, + path_or_buffer: PathOrBuffer | NoneType = None, + progress: ProgressCallback | None = None, ): """Maps the index from a file or buffer without loading it into memory. @@ -1136,7 +1177,9 @@ def view( :type progress: Optional[ProgressCallback], optional :raises Exception: If no source is defined. """ - assert not progress or _match_signature(progress, [int, int], bool), "Invalid callback signature" + assert not progress or _match_signature(progress, [int, int], bool), ( + "Invalid callback signature" + ) path_or_buffer = path_or_buffer if path_or_buffer is not None else self.path if path_or_buffer is None: @@ -1183,8 +1226,8 @@ def join( other: Index, max_proposals: int = 0, exact: bool = False, - progress: Optional[ProgressCallback] = None, - ) -> Dict[Key, Key]: + progress: ProgressCallback | None = None, + ) -> dict[Key, Key]: """Performs "Semantic Join" or pairwise matching between `self` & `other` index. Is different from `search`, as no collisions are allowed in resulting pairs. Uses the concept of "Stable Marriages" from Combinatorics, famous for the 2012 @@ -1201,7 +1244,9 @@ def join( :return: Mapping from keys of `self` to keys of `other` :rtype: Dict[Key, Key] """ - assert not progress or _match_signature(progress, [int, int], bool), "Invalid callback signature" + assert not progress or _match_signature(progress, [int, int], bool), ( + "Invalid callback signature" + ) return self._compiled.join( other=other._compiled, @@ -1213,13 +1258,13 @@ def join( def cluster( self, *, - vectors: Optional[np.ndarray] = None, - keys: Optional[np.ndarray] = None, - min_count: Optional[int] = None, - max_count: Optional[int] = None, + vectors: NDArray[Any] | None = None, + keys: NDArray[Any] | None = None, + min_count: int | None = None, + max_count: int | None = None, threads: int = 0, - log: Union[str, bool] = False, - progress: Optional[ProgressCallback] = None, + log: str | bool = False, + progress: ProgressCallback | None = None, ) -> Clustering: """ Clusters already indexed or provided `vectors`, mapping them to various centroids. @@ -1238,7 +1283,9 @@ def cluster( :return: Matches for one or more queries :rtype: Union[Matches, BatchMatches] """ - assert not progress or _match_signature(progress, [int, int], bool), "Invalid callback signature" + assert not progress or _match_signature(progress, [int, int], bool), ( + "Invalid callback signature" + ) if min_count is None: min_count = 0 @@ -1271,7 +1318,9 @@ def cluster( batch_matches = BatchMatches(*results) return Clustering(self, batch_matches, keys) - def pairwise_distance(self, left: KeyOrKeysLike, right: KeyOrKeysLike) -> Union[np.ndarray, float]: + def pairwise_distance( + self, left: KeyOrKeysLike, right: KeyOrKeysLike + ) -> NDArray[Any] | float: """Computes the pairwise distance between keys or key arrays. If `left` and `right` are single keys, returns the distance between them. @@ -1287,7 +1336,7 @@ def pairwise_distance(self, left: KeyOrKeysLike, right: KeyOrKeysLike) -> Union[ assert isinstance(left, Iterable) == isinstance(right, Iterable) if not isinstance(left, Iterable): - return self._compiled.pairwise_distance(int(left), int(right)) + return self._compiled.pairwise_distance(int(left), int(right)) # type: ignore[arg-type] else: left = np.array(left).astype(Key) right = np.array(right).astype(Key) @@ -1303,7 +1352,7 @@ def keys(self) -> IndexedKeys: return IndexedKeys(self) @property - def vectors(self) -> np.ndarray: + def vectors(self) -> Any: """Retrieves all vectors associated with the indexed keys. :return: Array of vectors. @@ -1354,7 +1403,7 @@ def stats(self) -> _CompiledIndexStats: return self._compiled.stats @property - def levels_stats(self) -> List[_CompiledIndexStats]: + def levels_stats(self) -> list[_CompiledIndexStats]: """Get the accumulated statistics for each level of the graph. :return: List of statistics for each level of the graph. @@ -1385,14 +1434,14 @@ def level_stats(self, level: int) -> _CompiledIndexStats: return self._compiled.level_stats(level) @property - def specs(self) -> Dict[str, Union[str, int, bool]]: + def specs(self) -> dict[str, str | int | bool]: """Returns the specifications of the index. :return: Dictionary of index specifications. :rtype: Dict[str, Union[str, int, bool]] """ if not hasattr(self, "_compiled"): - return "usearch.Index(failed)" + return {} return { "type": "usearch.Index", "ndim": self.ndim, @@ -1444,7 +1493,9 @@ def __repr_pretty__(self) -> str: """ if not hasattr(self, "_compiled"): return "usearch.Index(failed)" - level_stats = [f"--- {i}. {self.level_stats(i).nodes:,} nodes" for i in range(self.nlevels)] + level_stats = [ + f"--- {i}. {self.level_stats(i).nodes:,} nodes" for i in range(self.nlevels) + ] lines = "\n".join( [ "usearch.Index", @@ -1510,7 +1561,7 @@ def search( *, threads: int = 0, exact: bool = False, - progress: Optional[ProgressCallback] = None, + progress: ProgressCallback | None = None, ): return _search_in_compiled( self._compiled.search_many, @@ -1526,16 +1577,16 @@ def search( def search( - dataset: np.ndarray, - query: np.ndarray, + dataset: NDArray[Any], + query: NDArray[Any], count: int = 10, metric: MetricLike = MetricKind.Cos, *, exact: bool = False, threads: int = 0, - log: Union[str, bool] = False, - progress: Optional[ProgressCallback] = None, -) -> Union[Matches, BatchMatches]: + log: str | bool = False, + progress: ProgressCallback | None = None, +) -> Matches | BatchMatches: """Shortcut for search, that can avoid index construction. Particularly useful for tiny datasets, where brute-force exact search works fast enough. @@ -1564,7 +1615,9 @@ def search( :return: Matches for one or more queries :rtype: Union[Matches, BatchMatches] """ - assert not progress or _match_signature(progress, [int, int], bool), "Invalid callback signature" + assert not progress or _match_signature(progress, [int, int], bool), ( + "Invalid callback signature" + ) assert dataset.ndim == 2, "Dataset must be a matrix, with a vector in each row" if not exact: @@ -1635,8 +1688,8 @@ def kmeans( inertia_threshold: float = 1e-4, max_seconds: float = 60.0, min_shifts: float = 0.01, - seed: Optional[int] = None, -) -> Tuple[np.ndarray, np.ndarray, np.ndarray]: + seed: int | None = None, +) -> tuple[NDArray[Any], NDArray[Any], NDArray[Any]]: """ Performs KMeans clustering on a dataset using the USearch library with mixed-precision support. @@ -1706,7 +1759,11 @@ def kmeans( dtype = _normalize_dtype(dtype, ndim=X.shape[1], metric=metric) # Generating a 64-bit unsigned integer in NumPy may be somewhat tricky. - seed = np.random.default_rng().integers(0, 2**64, dtype=np.uint64) if seed is None else seed + actual_seed: int | np.unsignedinteger = ( + np.random.default_rng().integers(0, 2**64, dtype=np.uint64) + if seed is None + else seed + ) assignments, distances, centroids = _kmeans( X, k, @@ -1716,6 +1773,6 @@ def kmeans( min_shifts=min_shifts, inertia_threshold=inertia_threshold, dtype=dtype, - seed=seed, + seed=actual_seed, ) return assignments, distances, centroids diff --git a/python/usearch/io.py b/python/usearch/io.py index 417e9d36e..8299d90d9 100644 --- a/python/usearch/io.py +++ b/python/usearch/io.py @@ -1,9 +1,15 @@ +from __future__ import annotations + import os import struct import typing +from typing import TYPE_CHECKING, Any import numpy as np +if TYPE_CHECKING: + from numpy.typing import NDArray + def numpy_scalar_size(dtype) -> int: return { @@ -21,7 +27,7 @@ def numpy_scalar_size(dtype) -> int: }[dtype] -def guess_numpy_dtype_from_filename(filename) -> typing.Optional[type]: +def guess_numpy_dtype_from_filename(filename: str) -> type | None: if filename.endswith(".fbin"): return np.float32 elif filename.endswith(".dbin"): @@ -45,10 +51,10 @@ def guess_numpy_dtype_from_filename(filename) -> typing.Optional[type]: def load_matrix( filename: str, start_row: int = 0, - count_rows: int = None, + count_rows: int | None = None, view: bool = False, - dtype: typing.Optional[type] = None, -) -> typing.Optional[np.ndarray]: + dtype: type | None = None, +) -> NDArray[Any] | None: """Read *.ibin, *.bbib, *.hbin, *.fbin, *.dbin, *.i8bin, *.i32bin files with matrices. :param filename: path to the matrix file @@ -70,18 +76,22 @@ def load_matrix( with open(filename, "rb") as f: rows, cols = np.fromfile(f, count=2, dtype=np.int32).astype(np.uint64) - + # Validate file size matches expected data size f.seek(0, 2) # Go to end file_size = f.tell() expected_size = 8 + (rows * cols * scalar_size) # Header + data - + if file_size != expected_size: if file_size < expected_size: - raise ValueError(f"File {filename} is truncated. Expected {expected_size:,} bytes, got {file_size:,} bytes") + raise ValueError( + f"File {filename} is truncated. Expected {expected_size:,} bytes, got {file_size:,} bytes" + ) else: - raise ValueError(f"File {filename} is larger than expected. Expected {expected_size:,} bytes, got {file_size:,} bytes") - + raise ValueError( + f"File {filename} is larger than expected. Expected {expected_size:,} bytes, got {file_size:,} bytes" + ) + f.seek(8) # Back to start of data rows = (rows - start_row) if count_rows is None else count_rows row_offset = start_row * scalar_size * cols @@ -103,7 +113,7 @@ def load_matrix( ).reshape(rows, cols) -def save_matrix(vectors: np.ndarray, filename: str): +def save_matrix(vectors: NDArray[Any], filename: str): """Write *.ibin, *.bbib, *.hbin, *.fbin, *.dbin, *.i8bin, *.i32bin, *.f32bin files with matrices. :param vectors: the matrix to serialize @@ -111,6 +121,7 @@ def save_matrix(vectors: np.ndarray, filename: str): :param filename: path to the matrix file :type filename: str """ + dtype: type | np.dtype[typing.Any] if filename.endswith(".fbin"): dtype = np.float32 elif filename.endswith(".dbin"): diff --git a/python/usearch/numba.py b/python/usearch/numba.py index 3c1a4e61c..cb2f999a6 100644 --- a/python/usearch/numba.py +++ b/python/usearch/numba.py @@ -4,7 +4,7 @@ # into the primary `Index` class, connecting USearch with Numba. from math import sqrt -from usearch.index import MetricKind, ScalarKind, MetricSignature, CompiledMetric +from usearch.index import CompiledMetric, MetricKind, MetricSignature, ScalarKind def jit( @@ -25,12 +25,20 @@ def jit( assert isinstance(metric, MetricKind) assert isinstance(dtype, ScalarKind) - from numba import cfunc, types, carray + from numba import carray, cfunc, types # type: ignore[import-not-found] - signature_i8args = types.float32(types.CPointer(types.int8), types.CPointer(types.int8)) - signature_f16args = types.float32(types.CPointer(types.float16), types.CPointer(types.float16)) - signature_f32args = types.float32(types.CPointer(types.float32), types.CPointer(types.float32)) - signature_f64args = types.float32(types.CPointer(types.float64), types.CPointer(types.float64)) + signature_i8args = types.float32( + types.CPointer(types.int8), types.CPointer(types.int8) + ) + signature_f16args = types.float32( + types.CPointer(types.float16), types.CPointer(types.float16) + ) + signature_f32args = types.float32( + types.CPointer(types.float32), types.CPointer(types.float32) + ) + signature_f64args = types.float32( + types.CPointer(types.float64), types.CPointer(types.float64) + ) numba_supported_types = ( ScalarKind.I8, diff --git a/python/usearch/server.py b/python/usearch/server.py index 8f93ff6ac..96e2a6f37 100644 --- a/python/usearch/server.py +++ b/python/usearch/server.py @@ -1,16 +1,22 @@ #!/usr/bin/env python3 -# -*- coding: utf-8 -*- +from __future__ import annotations -import os import argparse +import os +from typing import TYPE_CHECKING, Any + import numpy as np -from typing import List +from ucall.rich_posix import Server # type: ignore[import-untyped] -from ucall.rich_posix import Server -from usearch.index import Index, Matches, Key +from usearch.index import Index, Key +if TYPE_CHECKING: + from numpy.typing import NDArray -def _ascii_to_vector(string: str) -> np.ndarray: + from usearch.index import BatchMatches, Matches + + +def _ascii_to_vector(string: str) -> NDArray[Any]: """ WARNING: A dirty performance hack! Assuming the `i8` vectors in our implementations are just integers, @@ -52,33 +58,31 @@ def ndim() -> int: @server def capacity() -> int: - return index.capacity() + return index.capacity @server def connectivity() -> int: - return index.connectivity() + return index.connectivity @server - def add_one(key: int, vector: np.ndarray): - print("adding", key, vector) + def add_one(key: int, vector: NDArray[Any]): keys = np.array([key], dtype=Key) vectors = vector.flatten().reshape(vector.shape[0], 1) index.add(keys, vectors) @server - def add_many(keys: np.ndarray, vectors: np.ndarray): + def add_many(keys: NDArray[Any], vectors: NDArray[Any]): index.add(keys, vectors, threads=threads) @server - def search_one(vector: np.ndarray, count: int) -> List[dict]: - print("search", vector, count) + def search_one(vector: NDArray[Any], count: int) -> list[tuple]: vectors = vector.reshape(vector.shape[0], 1) - results: Matches = index.search(vectors, count) + results: Matches | BatchMatches = index.search(vectors, count) return results.to_list() @server - def search_many(vectors: np.ndarray, count: int) -> List[List[dict]]: - results: Matches = index.search(vectors, count) + def search_many(vectors: NDArray[Any], count: int) -> list[tuple]: + results: Matches | BatchMatches = index.search(vectors, count) return results.to_list() @server @@ -100,7 +104,9 @@ def search_ascii(string: str, count: int): parser = argparse.ArgumentParser() parser.add_argument("-v", "--verbose", help="log server activity") parser.add_argument("--ndim", type=int, help="dimensionality of the vectors") - parser.add_argument("--immutable", type=bool, default=False, help="the index can not be updated") + parser.add_argument( + "--immutable", type=bool, default=False, help="the index can not be updated" + ) parser.add_argument( "--metric", @@ -116,8 +122,12 @@ def search_ascii(string: str, count: int): default=8545, help="port to open for client connections", ) - parser.add_argument("-j", "--threads", type=int, default=1, help="number of CPU threads to use") - parser.add_argument("--path", type=str, default="index.usearch", help="where to store the index") + parser.add_argument( + "-j", "--threads", type=int, default=1, help="number of CPU threads to use" + ) + parser.add_argument( + "--path", type=str, default="index.usearch", help="where to store the index" + ) args = parser.parse_args() assert args.ndim is not None, "Define the number of dimensions!" diff --git a/setup.py b/setup.py index 929c8b528..107954c93 100644 --- a/setup.py +++ b/setup.py @@ -159,6 +159,7 @@ def get_bool_env_w_name(name: str, preference: bool) -> tuple: install_requires = [ "numpy", "tqdm", + "ucall", ] if use_simsimd: include_dirs.append("simsimd/include") @@ -209,6 +210,11 @@ def get_bool_env_w_name(name: str, preference: bool) -> tuple: "Topic :: Database :: Database Engines/Servers", "Topic :: Scientific/Engineering :: Artificial Intelligence", ], + extras_require={ + "sklearn": ["scikit-learn>=1.0.0, <2.0.0",], + "graph": ["networkx"], + "plot": ["matplotlib"], + }, include_dirs=include_dirs, ext_modules=ext_modules, install_requires=install_requires, From 97ec51701fe32d43941c8d86b3b42eb6cead8706 Mon Sep 17 00:00:00 2001 From: Roman Solomatin <36135455+Samoed@users.noreply.github.com> Date: Sun, 5 Apr 2026 21:08:06 +0300 Subject: [PATCH 3/6] remove ucall --- setup.py | 1 - 1 file changed, 1 deletion(-) diff --git a/setup.py b/setup.py index cc7e620e0..8303fe89a 100644 --- a/setup.py +++ b/setup.py @@ -181,7 +181,6 @@ def get_bool_env_w_name(name: str, preference: bool) -> tuple: install_requires = [ "numpy", "tqdm", - "ucall", ] if use_numkong: include_dirs.append("numkong/include") From e2740b011a013ffe845cdb58f94044d390b2d8b2 Mon Sep 17 00:00:00 2001 From: Roman Solomatin <36135455+Samoed@users.noreply.github.com> Date: Mon, 6 Apr 2026 09:39:29 +0300 Subject: [PATCH 4/6] make test runnable --- python/usearch/__init__.py | 4 +++- python/usearch/index.py | 1 - 2 files changed, 3 insertions(+), 2 deletions(-) diff --git a/python/usearch/__init__.py b/python/usearch/__init__.py index cd18b736c..b25463c9e 100644 --- a/python/usearch/__init__.py +++ b/python/usearch/__init__.py @@ -30,11 +30,13 @@ from usearch.compiled import ( # type: ignore[import-not-found] + VERSION_MAJOR, + VERSION_MINOR, + VERSION_PATCH, # Default values: DEFAULT_CONNECTIVITY, DEFAULT_EXPANSION_ADD, DEFAULT_EXPANSION_SEARCH, - USES_FP16LIB, # Dependencies: USES_OPENMP, USES_NUMKONG, diff --git a/python/usearch/index.py b/python/usearch/index.py index a48c1f887..a392be890 100644 --- a/python/usearch/index.py +++ b/python/usearch/index.py @@ -21,7 +21,6 @@ DEFAULT_CONNECTIVITY, DEFAULT_EXPANSION_ADD, DEFAULT_EXPANSION_SEARCH, - USES_FP16LIB, USES_OPENMP, USES_SIMSIMD, MetricKind, From 7172af1b10420e69a952ec760774ccc5dfa032d2 Mon Sep 17 00:00:00 2001 From: Roman Solomatin <36135455+Samoed@users.noreply.github.com> Date: Mon, 6 Apr 2026 09:40:07 +0300 Subject: [PATCH 5/6] lint & format --- python/usearch/__init__.py | 28 ++++-------- python/usearch/client.py | 12 ++--- python/usearch/eval.py | 58 ++++++----------------- python/usearch/index.py | 94 +++++++++----------------------------- python/usearch/numba.py | 16 ++----- python/usearch/server.py | 12 ++--- 6 files changed, 56 insertions(+), 164 deletions(-) diff --git a/python/usearch/__init__.py b/python/usearch/__init__.py index b25463c9e..cc8964c26 100644 --- a/python/usearch/__init__.py +++ b/python/usearch/__init__.py @@ -30,22 +30,22 @@ from usearch.compiled import ( # type: ignore[import-not-found] - VERSION_MAJOR, - VERSION_MINOR, - VERSION_PATCH, # Default values: DEFAULT_CONNECTIVITY, DEFAULT_EXPANSION_ADD, DEFAULT_EXPANSION_SEARCH, + USES_NUMKONG, + USES_NUMKONG_DYNAMIC_DISPATCH, # Dependencies: USES_OPENMP, - USES_NUMKONG, USES_SIMSIMD, - USES_NUMKONG_DYNAMIC_DISPATCH, USES_SIMSIMD_DYNAMIC_DISPATCH, + VERSION_MAJOR, + VERSION_MINOR, + VERSION_PATCH, + hardware_acceleration_available, # Hardware capabilities: hardware_acceleration_compiled, - hardware_acceleration_available, ) __version__ = f"{VERSION_MAJOR}.{VERSION_MINOR}.{VERSION_PATCH}" @@ -88,9 +88,7 @@ def get_binary_name(self) -> tuple[str, str]: os_part = os_map.get(platform.system(), "") arch = platform.machine() arch_part = arch_map.get(arch, "") - extension = {"Linux": "so", "Windows": "dll", "Darwin": "dylib"}.get( - platform.system(), "" - ) + extension = {"Linux": "so", "Windows": "dll", "Darwin": "dylib"}.get(platform.system(), "") source_filename = f"usearch_sqlite_{os_part}_{arch_part}_{version}.{extension}" target_filename = f"usearch_sqlite.{extension}" return source_filename, target_filename @@ -125,9 +123,7 @@ def sqlite_found_or_downloaded(self) -> str | None: local_path = os.path.join(download_dir, target_filename) if not os.path.exists(local_path): # If not found locally, warn the user and download from GitHub - warnings.warn( - "Will download `usearch_sqlite` binary from GitHub.", UserWarning - ) + warnings.warn("Will download `usearch_sqlite` binary from GitHub.", UserWarning) try: source_url = self.determine_download_url(self.version, source_filename) os.makedirs(download_dir, exist_ok=True) @@ -135,9 +131,7 @@ def sqlite_found_or_downloaded(self) -> str | None: except HTTPError as e: # If the download fails due to HTTPError (e.g., 404 Not Found), like a missing lib version if e.code == 404: - warnings.warn( - f"Download failed: {e.url} could not be found.", UserWarning - ) + warnings.warn(f"Download failed: {e.url} could not be found.", UserWarning) else: warnings.warn( f"Download failed with HTTP error: {e.code} {e.reason}", @@ -150,9 +144,7 @@ def sqlite_found_or_downloaded(self) -> str | None: path_wout_extension, _, _ = local_path.rpartition(".") return path_wout_extension else: - warnings.warn( - "Failed to download `usearch_sqlite` binary from GitHub.", UserWarning - ) + warnings.warn("Failed to download `usearch_sqlite` binary from GitHub.", UserWarning) return None diff --git a/python/usearch/client.py b/python/usearch/client.py index 2c52bca3c..d3d677ac5 100644 --- a/python/usearch/client.py +++ b/python/usearch/client.py @@ -26,9 +26,7 @@ def _vector_to_ascii(vector: NDArray[Any]) -> str | None: class IndexClient: - def __init__( - self, uri: str = "127.0.0.1", port: int = 8545, use_http: bool = True - ) -> None: + def __init__(self, uri: str = "127.0.0.1", port: int = 8545, use_http: bool = True) -> None: self.client = Client(uri=uri, port=port, use_http=use_http) def add_one(self, key: int, vector: NDArray[Any]): @@ -50,9 +48,7 @@ def add_many(self, keys: NDArray[Any], vectors: NDArray[Any]): def add(self, keys: NDArray[Any] | int, vectors: NDArray[Any]): if isinstance(keys, int) or len(keys) == 1: - return self.add_one( - int(keys) if isinstance(keys, np.ndarray) else keys, vectors - ) + return self.add_one(int(keys) if isinstance(keys, np.ndarray) else keys, vectors) else: return self.add_many(keys, vectors) @@ -76,9 +72,7 @@ def search_one(self, vector: NDArray[Any], count: int) -> Matches: def search_many(self, vectors: NDArray[Any], count: int) -> BatchMatches: batch_size: int = vectors.shape[0] - list_of_matches: list[list[dict]] = self.client.search_many( - vectors=vectors, count=count - ) + list_of_matches: list[list[dict]] = self.client.search_many(vectors=vectors, count=count) keys = np.zeros((batch_size, count), dtype=np.uint32) distances = np.zeros((batch_size, count), dtype=np.float32) diff --git a/python/usearch/eval.py b/python/usearch/eval.py index f943b9523..417bbfbd2 100644 --- a/python/usearch/eval.py +++ b/python/usearch/eval.py @@ -94,9 +94,7 @@ class SearchStats: @property def mean_efficiency(self) -> float: - return 1 - float(self.computed_distances) / ( - self.count_queries * self.index_size - ) + return 1 - float(self.computed_distances) / (self.count_queries * self.index_size) @property def mean_recall(self) -> float: @@ -143,9 +141,7 @@ def self_recall(index: Index, sample: float | int = 1.0, **kwargs) -> SearchStat matches = index.search(vectors, **kwargs) count_matches: int = ( - matches.count_matches(keys) - if isinstance(matches, BatchMatches) - else int(matches.keys[0] == keys[0]) + matches.count_matches(keys) if isinstance(matches, BatchMatches) else int(matches.keys[0] == keys[0]) ) return SearchStats( index_size=len(index), @@ -210,9 +206,7 @@ def ndcg(relevances: NDArray[Any], k: int | None = None) -> float: return float(dcg(relevances, k) / best_dcg) -def relevance( - expected: NDArray[Any], predicted: NDArray[Any], k: int | None = None -) -> list[int]: +def relevance(expected: NDArray[Any], predicted: NDArray[Any], k: int | None = None) -> list[int]: """Calculate relevance scores. Binary relevance scores :param expected: ground-truth keys @@ -274,28 +268,20 @@ def build( assert loaded_vectors is not None, f"Failed to load vectors from {vectors}" d.vectors = loaded_vectors ndim = d.vectors.shape[1] - count = ( - min(d.vectors.shape[0], count) - if count is not None - else d.vectors.shape[0] - ) + count = min(d.vectors.shape[0], count) if count is not None else d.vectors.shape[0] d.vectors = d.vectors[:count, :] d.keys = np.arange(count, dtype=Key) if queries is not None: loaded_queries = load_matrix(queries) - assert loaded_queries is not None, ( - f"Failed to load queries from {queries}" - ) + assert loaded_queries is not None, f"Failed to load queries from {queries}" d.queries = loaded_queries else: d.queries = d.vectors if neighbors is not None: loaded_neighbors = load_matrix(neighbors) - assert loaded_neighbors is not None, ( - f"Failed to load neighbors from {neighbors}" - ) + assert loaded_neighbors is not None, f"Failed to load neighbors from {neighbors}" d.neighbors = loaded_neighbors if k is not None: d.neighbors = d.neighbors[:, :k] @@ -349,9 +335,7 @@ def __add__(self, other: TaskResult): result = TaskResult() if self.add_operations and other.add_operations: result.add_operations = self.add_operations + other.add_operations - result.add_per_second = result.add_operations / ( - self.add_seconds + other.add_seconds - ) + result.add_per_second = result.add_operations / (self.add_seconds + other.add_seconds) else: base = self if self.add_operations else other result.add_operations = base.add_operations @@ -363,9 +347,7 @@ def __add__(self, other: TaskResult): (self.recall_at_one or 0.0) * self.search_operations + (other.recall_at_one or 0.0) * other.search_operations ) / (self.search_operations + other.search_operations) - result.search_per_second = result.search_operations / ( - self.search_seconds + other.search_seconds - ) + result.search_per_second = result.search_operations / (self.search_seconds + other.search_seconds) else: base = self if self.search_operations else other result.search_operations = base.search_operations @@ -448,9 +430,7 @@ class SearchTask: neighbors: NDArray[Any] def __call__(self, index: Index) -> TaskResult: - dt, results = measure_seconds( - lambda: index.search(self.queries, self.neighbors.shape[1]) - ) + dt, results = measure_seconds(lambda: index.search(self.queries, self.neighbors.shape[1])) return TaskResult( search_per_second=self.queries.shape[0] / dt, @@ -476,9 +456,7 @@ class Evaluation: ndim: int @staticmethod - def for_dataset( - dataset: Dataset, batch_size: int = 0, clusters: int = 1 - ) -> Evaluation: + def for_dataset(dataset: Dataset, batch_size: int = 0, clusters: int = 1) -> Evaluation: tasks: list[AddTask | SearchTask] = [] assert dataset.vectors is not None and dataset.keys is not None assert dataset.queries is not None and dataset.neighbors is not None @@ -522,9 +500,7 @@ def __call__(self, index: Index, post_clean: bool = True) -> dict: import argparse # Initialize the argument parser - parser = argparse.ArgumentParser( - description="Evaluate vector search index for speed and accuracy." - ) + parser = argparse.ArgumentParser(description="Evaluate vector search index for speed and accuracy.") # Define expected arguments parser.add_argument( @@ -553,18 +529,14 @@ def __call__(self, index: Index, post_clean: bool = True) -> dict: ) parser.add_argument("--metric", type=str, required=False, help="Distance function.") parser.add_argument("--count", type=int, help="Number of vectors to use.") - parser.add_argument( - "--ndim", type=int, help="Number of dimensions for the vectors." - ) + parser.add_argument("--ndim", type=int, help="Number of dimensions for the vectors.") parser.add_argument( "--batch_size", type=int, default=0, help="Batch size for indexing and searching.", ) - parser.add_argument( - "--clusters", type=int, default=1, help="Number of clusters for indexing." - ) + parser.add_argument("--clusters", type=int, default=1, help="Number of clusters for indexing.") # Parse arguments from the command line args = parser.parse_args() @@ -579,9 +551,7 @@ def __call__(self, index: Index, post_clean: bool = True) -> dict: ) # Prepare the evaluation - evaluation = Evaluation.for_dataset( - dataset, batch_size=args.batch_size, clusters=args.clusters - ) + evaluation = Evaluation.for_dataset(dataset, batch_size=args.batch_size, clusters=args.clusters) index = Index(ndim=dataset.ndim, dtype=args.dtype, metric=args.metric) # Perform the evaluation diff --git a/python/usearch/index.py b/python/usearch/index.py index a392be890..b8e2d9004 100644 --- a/python/usearch/index.py +++ b/python/usearch/index.py @@ -16,11 +16,13 @@ from numpy.typing import NDArray from tqdm import tqdm +# Precompiled symbols that will be exposed # Precompiled symbols that will be exposed from usearch.compiled import ( # type: ignore[import-not-found] DEFAULT_CONNECTIVITY, DEFAULT_EXPANSION_ADD, DEFAULT_EXPANSION_SEARCH, + USES_NUMKONG, USES_OPENMP, USES_SIMSIMD, MetricKind, @@ -50,18 +52,7 @@ from usearch.compiled import ( kmeans as _kmeans, ) -# Precompiled symbols that will be exposed -from usearch.compiled import ( - MetricKind, - ScalarKind, - MetricSignature, - DEFAULT_CONNECTIVITY, - DEFAULT_EXPANSION_ADD, - DEFAULT_EXPANSION_SEARCH, - USES_OPENMP, - USES_NUMKONG, - USES_SIMSIMD, -) + if TYPE_CHECKING: from usearch.compiled import ( IndexStats as _CompiledIndexStats, # type: ignore[import-not-found] @@ -84,9 +75,7 @@ class CompiledMetric(NamedTuple): NoneType: TypeAlias = type(None) -KeyOrKeysLike: TypeAlias = ( - Key | Iterable[Key] | int | Iterable[int] | NDArray[Any] | memoryview -) +KeyOrKeysLike: TypeAlias = Key | Iterable[Key] | int | Iterable[int] | NDArray[Any] | memoryview VectorOrVectorsLike: TypeAlias = NDArray[Any] | Iterable[NDArray[Any]] | memoryview @@ -101,9 +90,7 @@ class CompiledMetric(NamedTuple): ProgressCallback = Callable[[int, int], bool] -def _match_signature( - func: Callable[..., Any], arg_types: list[type], ret_type: type -) -> bool: +def _match_signature(func: Callable[..., Any], arg_types: list[type], ret_type: type) -> bool: assert callable(func), "Not callable" sig = signature(func) param_types = [param.annotation for param in sig.parameters.values()] @@ -216,9 +203,7 @@ def _search_in_compiled( # assert isinstance(vectors, np.ndarray), "Expects a NumPy array" assert vectors.ndim == 1 or vectors.ndim == 2, "Expects a matrix or vector" - assert not progress or _match_signature(progress, [int, int], bool), ( - "Invalid callback" - ) + assert not progress or _match_signature(progress, [int, int], bool), "Invalid callback" if vectors.ndim == 1: vectors = vectors.reshape(1, len(vectors)) @@ -273,9 +258,7 @@ def _add_to_compiled( ) -> int | NDArray[Any]: # assert isinstance(vectors, np.ndarray), "Expects a NumPy array" - assert not progress or _match_signature(progress, [int, int], bool), ( - "Invalid callback" - ) + assert not progress or _match_signature(progress, [int, int], bool), "Invalid callback" assert vectors.ndim == 1 or vectors.ndim == 2, "Expects a matrix or vector" if vectors.ndim == 1: vectors = vectors.reshape(1, len(vectors)) @@ -360,10 +343,7 @@ def __iter__(self): def to_list(self) -> list[tuple]: """Convert to list of (key, distance) tuples.""" - return [ - (int(key), float(distance)) - for key, distance in zip(self.keys, self.distances) - ] + return [(int(key), float(distance)) for key, distance in zip(self.keys, self.distances)] def __repr__(self) -> str: return f"usearch.Matches({len(self)})" @@ -636,9 +616,7 @@ def __init__( self._metric_pointer = metric.pointer self._metric_signature = metric.signature else: - raise ValueError( - "The `metric` must be a `CompiledMetric` or a `MetricKind`" - ) + raise ValueError("The `metric` must be a `CompiledMetric` or a `MetricKind`") # Validate, that the right scalar type is defined dtype = _normalize_dtype(dtype, ndim, self._metric_kind) @@ -676,9 +654,7 @@ def metadata(path_or_buffer: PathOrBuffer) -> dict | None: raise e @staticmethod - def restore( - path_or_buffer: PathOrBuffer, view: bool = False, **kwargs - ) -> Index | None: + def restore(path_or_buffer: PathOrBuffer, view: bool = False, **kwargs) -> Index | None: meta = Index.metadata(path_or_buffer) if not meta: return None @@ -840,9 +816,7 @@ def get( dtype = _normalize_dtype(dtype) view_dtype = _to_numpy_dtype(dtype) if view_dtype is None: - raise NotImplementedError( - "The requested representation type is not supported by NumPy" - ) + raise NotImplementedError("The requested representation type is not supported by NumPy") def cast(result): if result is not None: @@ -858,11 +832,7 @@ def cast(result): actual_keys = np.array(list(keys), dtype=Key) # type: ignore[arg-type] results = self._compiled.get_many(actual_keys, dtype) - results = ( - cast(results) - if isinstance(results, np.ndarray) - else [cast(result) for result in results] - ) + results = cast(results) if isinstance(results, np.ndarray) else [cast(result) for result in results] return results[0] if is_one else results def __getitem__(self, keys: KeyOrKeysLike) -> Any: @@ -1025,9 +995,7 @@ def metric(self, metric: MetricLike): metric_pointer = metric.pointer metric_signature = metric.signature else: - raise ValueError( - "The `metric` must be a `CompiledMetric` or a `MetricKind`" - ) + raise ValueError("The `metric` must be a `CompiledMetric` or a `MetricKind`") return self._compiled.change_metric( metric_kind=metric_kind, @@ -1131,9 +1099,7 @@ def save( :return: The index data as bytes if saving to a buffer, otherwise None. :rtype: Optional[bytes] """ - assert not progress or _match_signature(progress, [int, int], bool), ( - "Invalid callback signature" - ) + assert not progress or _match_signature(progress, [int, int], bool), "Invalid callback signature" path_or_buffer = path_or_buffer if path_or_buffer is not None else self.path if path_or_buffer is None: @@ -1158,9 +1124,7 @@ def load( :raises Exception: If no source is defined. :raises RuntimeError: If the file does not exist. """ - assert not progress or _match_signature(progress, [int, int], bool), ( - "Invalid callback signature" - ) + assert not progress or _match_signature(progress, [int, int], bool), "Invalid callback signature" path_or_buffer = path_or_buffer if path_or_buffer is not None else self.path if path_or_buffer is None: @@ -1189,9 +1153,7 @@ def view( :type progress: Optional[ProgressCallback], optional :raises Exception: If no source is defined. """ - assert not progress or _match_signature(progress, [int, int], bool), ( - "Invalid callback signature" - ) + assert not progress or _match_signature(progress, [int, int], bool), "Invalid callback signature" path_or_buffer = path_or_buffer if path_or_buffer is not None else self.path if path_or_buffer is None: @@ -1256,9 +1218,7 @@ def join( :return: Mapping from keys of `self` to keys of `other` :rtype: Dict[Key, Key] """ - assert not progress or _match_signature(progress, [int, int], bool), ( - "Invalid callback signature" - ) + assert not progress or _match_signature(progress, [int, int], bool), "Invalid callback signature" return self._compiled.join( other=other._compiled, @@ -1295,9 +1255,7 @@ def cluster( :return: Matches for one or more queries :rtype: Union[Matches, BatchMatches] """ - assert not progress or _match_signature(progress, [int, int], bool), ( - "Invalid callback signature" - ) + assert not progress or _match_signature(progress, [int, int], bool), "Invalid callback signature" if min_count is None: min_count = 0 @@ -1330,9 +1288,7 @@ def cluster( batch_matches = BatchMatches(*results) return Clustering(self, batch_matches, keys) - def pairwise_distance( - self, left: KeyOrKeysLike, right: KeyOrKeysLike - ) -> NDArray[Any] | float: + def pairwise_distance(self, left: KeyOrKeysLike, right: KeyOrKeysLike) -> NDArray[Any] | float: """Computes the pairwise distance between keys or key arrays. If `left` and `right` are single keys, returns the distance between them. @@ -1505,9 +1461,7 @@ def __repr_pretty__(self) -> str: """ if not hasattr(self, "_compiled"): return "usearch.Index(failed)" - level_stats = [ - f"--- {i}. {self.level_stats(i).nodes:,} nodes" for i in range(self.nlevels) - ] + level_stats = [f"--- {i}. {self.level_stats(i).nodes:,} nodes" for i in range(self.nlevels)] lines = "\n".join( [ "usearch.Index", @@ -1626,9 +1580,7 @@ def search( :return: Matches for one or more queries :rtype: Union[Matches, BatchMatches] """ - assert not progress or _match_signature(progress, [int, int], bool), ( - "Invalid callback signature" - ) + assert not progress or _match_signature(progress, [int, int], bool), "Invalid callback signature" assert dataset.ndim == 2, "Dataset must be a matrix, with a vector in each row" if not exact: @@ -1771,9 +1723,7 @@ def kmeans( # Generating a 64-bit unsigned integer in NumPy may be somewhat tricky. actual_seed: int | np.unsignedinteger = ( - np.random.default_rng().integers(0, 2**64, dtype=np.uint64) - if seed is None - else seed + np.random.default_rng().integers(0, 2**64, dtype=np.uint64) if seed is None else seed ) assignments, distances, centroids = _kmeans( X, diff --git a/python/usearch/numba.py b/python/usearch/numba.py index cb2f999a6..40d53a52a 100644 --- a/python/usearch/numba.py +++ b/python/usearch/numba.py @@ -27,18 +27,10 @@ def jit( from numba import carray, cfunc, types # type: ignore[import-not-found] - signature_i8args = types.float32( - types.CPointer(types.int8), types.CPointer(types.int8) - ) - signature_f16args = types.float32( - types.CPointer(types.float16), types.CPointer(types.float16) - ) - signature_f32args = types.float32( - types.CPointer(types.float32), types.CPointer(types.float32) - ) - signature_f64args = types.float32( - types.CPointer(types.float64), types.CPointer(types.float64) - ) + signature_i8args = types.float32(types.CPointer(types.int8), types.CPointer(types.int8)) + signature_f16args = types.float32(types.CPointer(types.float16), types.CPointer(types.float16)) + signature_f32args = types.float32(types.CPointer(types.float32), types.CPointer(types.float32)) + signature_f64args = types.float32(types.CPointer(types.float64), types.CPointer(types.float64)) numba_supported_types = ( ScalarKind.I8, diff --git a/python/usearch/server.py b/python/usearch/server.py index 96e2a6f37..d30e94a1e 100644 --- a/python/usearch/server.py +++ b/python/usearch/server.py @@ -104,9 +104,7 @@ def search_ascii(string: str, count: int): parser = argparse.ArgumentParser() parser.add_argument("-v", "--verbose", help="log server activity") parser.add_argument("--ndim", type=int, help="dimensionality of the vectors") - parser.add_argument( - "--immutable", type=bool, default=False, help="the index can not be updated" - ) + parser.add_argument("--immutable", type=bool, default=False, help="the index can not be updated") parser.add_argument( "--metric", @@ -122,12 +120,8 @@ def search_ascii(string: str, count: int): default=8545, help="port to open for client connections", ) - parser.add_argument( - "-j", "--threads", type=int, default=1, help="number of CPU threads to use" - ) - parser.add_argument( - "--path", type=str, default="index.usearch", help="where to store the index" - ) + parser.add_argument("-j", "--threads", type=int, default=1, help="number of CPU threads to use") + parser.add_argument("--path", type=str, default="index.usearch", help="where to store the index") args = parser.parse_args() assert args.ndim is not None, "Define the number of dimensions!" From 5408f87a3b0f9e9e32c38ae37af49613c37020ce Mon Sep 17 00:00:00 2001 From: Roman Solomatin <36135455+Samoed@users.noreply.github.com> Date: Mon, 6 Apr 2026 10:35:53 +0300 Subject: [PATCH 6/6] strict typing --- pyproject.toml | 1 + python/usearch/__init__.py | 2 +- python/usearch/client.py | 28 ++--- python/usearch/eval.py | 33 ++--- python/usearch/index.py | 239 ++++++++++++++++++++----------------- python/usearch/io.py | 4 +- python/usearch/numba.py | 9 +- python/usearch/server.py | 42 +++---- 8 files changed, 188 insertions(+), 170 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index dbaf6a154..db3bd3781 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -124,6 +124,7 @@ future-annotations = true # For TC rules [tool.mypy] python_version = "3.10" +strict = true [[tool.mypy.overrides]] # these modules not typed and don't have stubs diff --git a/python/usearch/__init__.py b/python/usearch/__init__.py index cc8964c26..3f5d62922 100644 --- a/python/usearch/__init__.py +++ b/python/usearch/__init__.py @@ -58,7 +58,7 @@ def __init__(self, version: str | None = None): self.version = version or __version__ @staticmethod - def determine_download_dir(): + def determine_download_dir() -> str: # Check if running within a virtual environment virtual_env = os.getenv("VIRTUAL_ENV") if virtual_env: diff --git a/python/usearch/client.py b/python/usearch/client.py index d3d677ac5..b5fb9bd2b 100644 --- a/python/usearch/client.py +++ b/python/usearch/client.py @@ -3,7 +3,7 @@ from typing import TYPE_CHECKING, Any import numpy as np -from ucall.client import Client # type: ignore[import-untyped] +from ucall.client import Client # type: ignore[import-not-found] if TYPE_CHECKING: from numpy.typing import NDArray @@ -29,7 +29,7 @@ class IndexClient: def __init__(self, uri: str = "127.0.0.1", port: int = 8545, use_http: bool = True) -> None: self.client = Client(uri=uri, port=port, use_http=use_http) - def add_one(self, key: int, vector: NDArray[Any]): + def add_one(self, key: int, vector: NDArray[Any]) -> None: assert isinstance(key, int) assert isinstance(vector, np.ndarray) vector = vector.flatten() @@ -39,14 +39,14 @@ def add_one(self, key: int, vector: NDArray[Any]): else: self.client.add_one(key=key, vectors=vector) - def add_many(self, keys: NDArray[Any], vectors: NDArray[Any]): + def add_many(self, keys: NDArray[Any], vectors: NDArray[Any]) -> None: assert isinstance(keys, int) assert isinstance(vectors, np.ndarray) assert keys.ndim == 1 and vectors.ndim == 2 assert keys.shape[0] == vectors.shape[0] self.client.add_many(keys=keys, vectors=vectors) - def add(self, keys: NDArray[Any] | int, vectors: NDArray[Any]): + def add(self, keys: NDArray[Any] | int, vectors: NDArray[Any]) -> None: if isinstance(keys, int) or len(keys) == 1: return self.add_one(int(keys) if isinstance(keys, np.ndarray) else keys, vectors) else: @@ -60,7 +60,7 @@ def search_one(self, vector: NDArray[Any], count: int) -> Matches: else: raw = self.client.search_one(vector=vector, count=count) - matches: list[dict] = raw.json + matches: list[dict[str, Any]] = raw.json keys = np.array(count, dtype=np.uint32) distances = np.array(count, dtype=np.float32) @@ -72,7 +72,7 @@ def search_one(self, vector: NDArray[Any], count: int) -> Matches: def search_many(self, vectors: NDArray[Any], count: int) -> BatchMatches: batch_size: int = vectors.shape[0] - list_of_matches: list[list[dict]] = self.client.search_many(vectors=vectors, count=count) + list_of_matches: list[list[dict[str, Any]]] = self.client.search_many(vectors=vectors, count=count) keys = np.zeros((batch_size, count), dtype=np.uint32) distances = np.zeros((batch_size, count), dtype=np.float32) @@ -91,26 +91,26 @@ def search(self, vectors: NDArray[Any], count: int) -> Matches | BatchMatches: else: return self.search_many(vectors, count) - def __len__(self): - return self.client.size().json() + def __len__(self) -> int: + return int(self.client.size().json()) @property - def ndim(self): + def ndim(self) -> Any: return self.client.ndim().json() - def capacity(self): + def capacity(self) -> Any: return self.client.capacity().json() - def connectivity(self): + def connectivity(self) -> Any: return self.client.connectivity().json() - def load(self, path: str): + def load(self, path: str) -> None: raise NotImplementedError() - def view(self, path: str): + def view(self, path: str) -> None: raise NotImplementedError() - def save(self, path: str): + def save(self, path: str) -> None: raise NotImplementedError() diff --git a/python/usearch/eval.py b/python/usearch/eval.py index 417bbfbd2..a5bd84745 100644 --- a/python/usearch/eval.py +++ b/python/usearch/eval.py @@ -4,7 +4,7 @@ from dataclasses import asdict, dataclass from math import ceil from time import time_ns -from typing import TYPE_CHECKING, Any +from typing import TYPE_CHECKING, Any, cast import numpy as np @@ -65,7 +65,7 @@ def random_vectors( else: x = x.astype(_to_numpy_dtype(dtype)) if metric == MetricKind.IP: - return x / np.linalg.norm(x, axis=1, keepdims=True) + return cast("NDArray[Any]", x / np.linalg.norm(x, axis=1, keepdims=True)) return x @@ -101,7 +101,7 @@ def mean_recall(self) -> float: return self.count_matches / self.count_queries -def self_recall(index: Index, sample: float | int = 1.0, **kwargs) -> SearchStats: +def self_recall(index: Index, sample: float | int = 1.0, **kwargs: Any) -> SearchStats: """Simplest benchmark for a quality of search, which queries every existing member of the index, to make sure approximate search finds the point itself. @@ -152,7 +152,7 @@ def self_recall(index: Index, sample: float | int = 1.0, **kwargs) -> SearchStat ) -def measure_seconds(f: Callable) -> tuple[float, Any]: +def measure_seconds(f: Callable[[], Any]) -> tuple[float, Any]: """Simple function profiling decorator. :param f: Function to be profiled @@ -226,13 +226,14 @@ class Dataset: queries: NDArray[Any] | None neighbors: NDArray[Any] | None - def crop_neighbors(self, k: int): + def crop_neighbors(self, k: int) -> None: assert self.neighbors is not None self.neighbors = self.neighbors[:, k] @property - def ndim(self): - return self.vectors.shape[1] + def ndim(self) -> int: + assert self.vectors is not None + return int(self.vectors.shape[1]) @staticmethod def build( @@ -242,7 +243,7 @@ def build( count: int | None = None, ndim: int | None = None, k: int | None = None, - ): + ) -> Dataset: """Either loads an existing dataset from disk, or generates one on the fly. :param vectors: _description_, defaults to None @@ -331,7 +332,7 @@ def search_seconds(self) -> float: assert self.search_operations is not None and self.search_per_second is not None return self.search_operations / self.search_per_second - def __add__(self, other: TaskResult): + def __add__(self, other: TaskResult) -> TaskResult: result = TaskResult() if self.add_operations and other.add_operations: result.add_operations = self.add_operations + other.add_operations @@ -374,14 +375,14 @@ def __call__(self, index: Index) -> TaskResult: ) @property - def ndim(self): - return self.vectors.shape[1] + def ndim(self) -> int: + return int(self.vectors.shape[1]) @property - def count(self): - return self.vectors.shape[0] + def count(self) -> int: + return int(self.vectors.shape[0]) - def inplace_shuffle(self): + def inplace_shuffle(self) -> None: """Reorders the `vectors` and `keys`. Often used for robustness benchmarks.""" new_order = np.arange(self.count) @@ -403,7 +404,7 @@ def slices(self, batch_size: int) -> list[AddTask]: def clusters(self, number_of_clusters: int) -> list[AddTask]: """Splits this dataset into smaller chunks.""" - from sklearn.cluster import KMeans # type: ignore[import-untyped] + from sklearn.cluster import KMeans # type: ignore[import-not-found] clustering = KMeans( n_clusters=number_of_clusters, @@ -479,7 +480,7 @@ def for_dataset(dataset: Dataset, batch_size: int = 0, clusters: int = 1) -> Eva ndim=add.ndim, ) - def __call__(self, index: Index, post_clean: bool = True) -> dict: + def __call__(self, index: Index, post_clean: bool = True) -> dict[str, Any]: task_result = TaskResult() try: diff --git a/python/usearch/index.py b/python/usearch/index.py index b8e2d9004..88ff0f44d 100644 --- a/python/usearch/index.py +++ b/python/usearch/index.py @@ -7,10 +7,10 @@ # Python tooling, linters, and static analyzers. It also embeds JIT # into the primary `Index` class, connecting USearch with Numba. import os -from collections.abc import Callable, Iterable, Sequence +from collections.abc import Callable, Iterable, Iterator, Sequence from dataclasses import dataclass from inspect import signature -from typing import TYPE_CHECKING, Any, NamedTuple, TypeAlias, overload +from typing import TYPE_CHECKING, Any, NamedTuple, Protocol, TypeAlias, cast, overload import numpy as np from numpy.typing import NDArray @@ -25,9 +25,6 @@ USES_NUMKONG, USES_OPENMP, USES_SIMSIMD, - MetricKind, - MetricSignature, - ScalarKind, ) # Precompiled symbols that won't be exposed directly: @@ -37,6 +34,15 @@ from usearch.compiled import ( Indexes as _CompiledIndexes, ) +from usearch.compiled import ( + MetricKind as MetricKind, +) +from usearch.compiled import ( + MetricSignature as MetricSignature, +) +from usearch.compiled import ( + ScalarKind as ScalarKind, +) from usearch.compiled import ( exact_search as _exact_search, ) @@ -55,7 +61,7 @@ if TYPE_CHECKING: from usearch.compiled import ( - IndexStats as _CompiledIndexStats, # type: ignore[import-not-found] + IndexStats as _CompiledIndexStats, ) MetricKindBitwise = ( @@ -85,7 +91,7 @@ class CompiledMetric(NamedTuple): BytesLike: TypeAlias = bytes | bytearray | memoryview -PathOrBuffer: TypeAlias = str | os.PathLike | BytesLike +PathOrBuffer: TypeAlias = str | os.PathLike[str] | BytesLike ProgressCallback = Callable[[int, int], bool] @@ -98,7 +104,7 @@ def _match_signature(func: Callable[..., Any], arg_types: list[type], ret_type: def _normalize_dtype( - dtype, + dtype: ScalarKind | str | type | None, ndim: int = 0, metric: MetricKind = MetricKind.Cos, ) -> ScalarKind: @@ -142,22 +148,20 @@ def _normalize_dtype( return _normalize[dtype] -def _to_numpy_dtype(dtype: ScalarKind): +def _to_numpy_dtype(dtype: ScalarKind) -> type | None: if dtype == ScalarKind.BF16: return None - _normalize = { + _normalize: dict[ScalarKind, type] = { ScalarKind.F64: np.float64, ScalarKind.F32: np.float32, ScalarKind.F16: np.float16, ScalarKind.I8: np.int8, ScalarKind.B1: np.uint8, } - if dtype in _normalize.values(): - return dtype return _normalize[dtype] -def _normalize_metric(metric) -> MetricKind: +def _normalize_metric(metric: str | MetricKind | CompiledMetric | None) -> MetricKind | CompiledMetric: if metric is None: return MetricKind.Cos @@ -192,13 +196,20 @@ def _is_buffer(obj: Any) -> bool: return False +_BatchSearchTuple: TypeAlias = tuple[NDArray[Any], NDArray[Any], NDArray[Any], int, int] + + +class _SearchCallable(Protocol): + def __call__(self, vectors: NDArray[Any], *, progress: ProgressCallback | None = None, **kwargs: Any) -> _BatchSearchTuple: ... + + def _search_in_compiled( - compiled_callable: Callable, + compiled_callable: _SearchCallable, vectors: VectorOrVectorsLike, *, log: str | bool, progress: ProgressCallback | None, - **kwargs, + **kwargs: Any, ) -> Matches | BatchMatches: # assert isinstance(vectors, np.ndarray), "Expects a NumPy array" @@ -247,10 +258,10 @@ def update_progress_bar(processed: int, total: int) -> bool: def _add_to_compiled( - compiled, + compiled: Any, *, - keys, - vectors, + keys: KeyOrKeysLike | None, + vectors: VectorOrVectorsLike, copy: bool, threads: int, log: str | bool, @@ -270,9 +281,10 @@ def _add_to_compiled( start_id = len(compiled) keys = np.arange(start_id, start_id + count_vectors, dtype=Key) else: + assert keys is not None if not isinstance(keys, Iterable): assert count_vectors == 1, "Each vector must have a key" - keys = [keys] + keys = cast("Iterable[int]", [keys]) keys = np.array(keys).astype(Key) assert len(keys) == count_vectors @@ -301,7 +313,7 @@ def update_progress_bar(processed: int, total: int) -> bool: else: compiled.add_many(keys, vectors, copy=copy, threads=threads, progress=progress) - return keys + return cast("NDArray[Any]", keys) @dataclass @@ -311,7 +323,7 @@ class Match: key: int distance: float - def to_tuple(self) -> tuple: + def to_tuple(self) -> tuple[int, float]: return self.key, self.distance @@ -337,11 +349,11 @@ def __getitem__(self, index: int) -> Match: else: raise IndexError(f"`index` must be an integer under {len(self)}") - def __iter__(self): + def __iter__(self) -> Iterator[Match]: for i in range(len(self)): yield self[i] - def to_list(self) -> list[tuple]: + def to_list(self) -> list[tuple[int, float]]: """Convert to list of (key, distance) tuples.""" return [(int(key), float(distance)) for key, distance in zip(self.keys, self.distances)] @@ -350,7 +362,7 @@ def __repr__(self) -> str: @dataclass -class BatchMatches(Sequence): +class BatchMatches(Sequence[Matches]): """Search results for multiple queries in batch operations. Unused positions in arrays contain sentinel values (default keys, max distances). @@ -393,7 +405,7 @@ def __getitem__(self, index: int | slice) -> Matches | list[Matches]: else: raise IndexError(f"`index` must be an integer under {len(self)}") - def to_list(self) -> list[tuple]: + def to_list(self) -> list[tuple[int, float]]: """Flatten matches for all queries into a list of `(key, distance)` tuples.""" list_of_matches = [self.__getitem__(row) for row in range(self.__len__())] return [match.to_tuple() for matches in list_of_matches for match in matches] @@ -445,13 +457,13 @@ def centroids_popularity(self) -> tuple[NDArray[Any], NDArray[Any]]: return np.unique(self.matches.keys, return_counts=True) def members_of(self, centroid: Key) -> NDArray[Any]: - return self.queries[self.matches.keys.flatten() == centroid] + return cast("NDArray[Any]", self.queries[self.matches.keys.flatten() == centroid]) - def subcluster(self, centroid: Key, **clustering_kwargs) -> Clustering: + def subcluster(self, centroid: Key, **clustering_kwargs: Any) -> Clustering: sub_keys = self.members_of(centroid) return self.index.cluster(keys=sub_keys, **clustering_kwargs) - def plot_centroids_popularity(self): + def plot_centroids_popularity(self) -> None: from matplotlib import pyplot as plt _, sizes = self.centroids_popularity @@ -460,12 +472,12 @@ def plot_centroids_popularity(self): plt.show() @property - def network(self): + def network(self) -> Any: import networkx as nx keys, sizes = self.centroids_popularity - g = nx.Graph() + g: Any = nx.Graph() for key, size in zip(keys, sizes): g.add_node(key, size=size) @@ -477,7 +489,7 @@ def network(self): return g -class IndexedKeys(Sequence): +class IndexedKeys(Sequence[Key]): """View of all keys in the index.""" def __init__(self, index: Index) -> None: @@ -494,11 +506,11 @@ def __getitem__( # type: ignore[override] start, stop, step = offset_offsets_or_slice.indices(len(self)) if step != 1: raise ValueError("Slicing with a step is not supported") - return self._idx._compiled.get_keys_in_slice(start, stop - start) + return cast("NDArray[Any]", self._idx._compiled.get_keys_in_slice(start, stop - start)) elif isinstance(offset_offsets_or_slice, Iterable): offsets = np.array(offset_offsets_or_slice) - return self._idx._compiled.get_keys_at_offsets(offsets) + return cast("NDArray[Any]", self._idx._compiled.get_keys_at_offsets(offsets)) else: offset = int(offset_offsets_or_slice) @@ -506,12 +518,12 @@ def __getitem__( # type: ignore[override] offset += len(self) if offset < 0 or offset >= len(self): raise IndexError("Index out of range") - return self._idx._compiled.get_key_at_offset(offset) + return cast("Key", self._idx._compiled.get_key_at_offset(offset)) - def __array__(self, dtype=None) -> NDArray[Any]: + def __array__(self, dtype: type | None = None) -> NDArray[Any]: if dtype is None: dtype = Key - return self._idx._compiled.get_keys_in_slice().astype(dtype) + return cast("NDArray[Any]", self._idx._compiled.get_keys_in_slice().astype(dtype)) class Index: @@ -537,7 +549,7 @@ def __init__( expansion_add: int | None = None, expansion_search: int | None = None, multi: bool = False, - path: os.PathLike | None = None, + path: os.PathLike[str] | None = None, view: bool = False, enable_key_lookups: bool = True, ) -> None: @@ -641,20 +653,20 @@ def __init__( self.load(path) @staticmethod - def metadata(path_or_buffer: PathOrBuffer) -> dict | None: + def metadata(path_or_buffer: PathOrBuffer) -> dict[str, Any] | None: try: if _is_buffer(path_or_buffer): - return _index_dense_metadata_from_buffer(path_or_buffer) + return cast("dict[str, Any]", _index_dense_metadata_from_buffer(path_or_buffer)) else: - path_or_buffer = os.fspath(path_or_buffer) - if not os.path.exists(path_or_buffer): + path = os.fspath(cast("str | os.PathLike[str]", path_or_buffer)) + if not os.path.exists(path): return None - return _index_dense_metadata_from_path(path_or_buffer) + return cast("dict[str, Any]", _index_dense_metadata_from_path(path)) except Exception as e: raise e @staticmethod - def restore(path_or_buffer: PathOrBuffer, view: bool = False, **kwargs) -> Index | None: + def restore(path_or_buffer: PathOrBuffer, view: bool = False, **kwargs: Any) -> Index | None: meta = Index.metadata(path_or_buffer) if not meta: return None @@ -673,7 +685,7 @@ def restore(path_or_buffer: PathOrBuffer, view: bool = False, **kwargs) -> Index return index def __len__(self) -> int: - return self._compiled.__len__() + return int(self._compiled.__len__()) def add( self, @@ -776,18 +788,18 @@ def search( def contains(self, keys: KeyOrKeysLike) -> bool | NDArray[Any]: if isinstance(keys, Iterable): - return self._compiled.contains_many(np.array(keys, dtype=Key)) + return cast("NDArray[Any]", self._compiled.contains_many(np.array(keys, dtype=Key))) else: - return self._compiled.contains_one(int(keys)) + return cast("bool", self._compiled.contains_one(int(keys))) def __contains__(self, keys: KeyOrKeysLike) -> bool | NDArray[Any]: return self.contains(keys) def count(self, keys: KeyOrKeysLike) -> int | NDArray[Any]: if isinstance(keys, Iterable): - return self._compiled.count_many(np.array(keys, dtype=Key)) + return cast("NDArray[Any]", self._compiled.count_many(np.array(keys, dtype=Key))) else: - return self._compiled.count_one(int(keys)) + return cast("int", self._compiled.count_one(int(keys))) def get( self, @@ -818,7 +830,7 @@ def get( if view_dtype is None: raise NotImplementedError("The requested representation type is not supported by NumPy") - def cast(result): + def _cast_result(result: NDArray[Any] | None) -> NDArray[Any] | None: if result is not None: return result.view(view_dtype) return result @@ -831,8 +843,8 @@ def cast(result): else: actual_keys = np.array(list(keys), dtype=Key) # type: ignore[arg-type] - results = self._compiled.get_many(actual_keys, dtype) - results = cast(results) if isinstance(results, np.ndarray) else [cast(result) for result in results] + raw = self._compiled.get_many(actual_keys, dtype) + results: Any = _cast_result(raw) if isinstance(raw, np.ndarray) else [_cast_result(result) for result in raw] return results[0] if is_one else results def __getitem__(self, keys: KeyOrKeysLike) -> Any: @@ -873,10 +885,10 @@ def remove( :type: Union[int, np.ndarray] """ if not isinstance(keys, Iterable): - return self._compiled.remove_one(keys, compact=compact, threads=threads) + return cast("int", self._compiled.remove_one(keys, compact=compact, threads=threads)) else: keys = np.array(keys, dtype=Key) - return self._compiled.remove_many(keys, compact=compact, threads=threads) + return cast("NDArray[Any]", self._compiled.remove_many(keys, compact=compact, threads=threads)) def __delitem__(self, keys: KeyOrKeysLike) -> int | NDArray[Any]: return self.remove(keys) @@ -903,13 +915,13 @@ def rename( from_ = np.array(from_, dtype=Key) if isinstance(to, Iterable): to = np.array(to, dtype=Key) - return self._compiled.rename_many_to_many(from_, to) + return cast("NDArray[Any]", self._compiled.rename_many_to_many(from_, to)) else: - return self._compiled.rename_many_to_one(from_, int(to)) + return cast("int", self._compiled.rename_many_to_one(from_, int(to))) else: - return self._compiled.rename_one_to_one(int(from_), int(to)) # type: ignore[arg-type] + return cast("int", self._compiled.rename_one_to_one(int(from_), int(to))) # type: ignore[arg-type] @property def jit(self) -> bool: @@ -930,7 +942,7 @@ def hardware_acceleration(self) -> str: :return: "auto" if no hardware acceleration is available, otherwise an ISA subset name. :rtype: str """ - return self._compiled.hardware_acceleration + return str(self._compiled.hardware_acceleration) @property def size(self) -> int: @@ -939,7 +951,7 @@ def size(self) -> int: :return: The number of vectors in the index. :rtype: int """ - return self._compiled.size + return int(self._compiled.size) @property def ndim(self) -> int: @@ -948,7 +960,7 @@ def ndim(self) -> int: :return: The dimensionality of vectors in the index. :rtype: int """ - return self._compiled.ndim + return int(self._compiled.ndim) @property def serialized_length(self) -> int: @@ -957,7 +969,7 @@ def serialized_length(self) -> int: :return: The serialized length of the index in bytes. :rtype: int """ - return self._compiled.serialized_length + return int(self._compiled.serialized_length) @property def metric_kind(self) -> MetricKind | CompiledMetric: @@ -978,7 +990,7 @@ def metric(self) -> MetricKind | CompiledMetric: return self._metric_jit if self._metric_jit else self._metric_kind @metric.setter - def metric(self, metric: MetricLike): + def metric(self, metric: MetricLike) -> None: """Sets a new metric for the index. :param metric: The new metric to be used. @@ -997,7 +1009,7 @@ def metric(self, metric: MetricLike): else: raise ValueError("The `metric` must be a `CompiledMetric` or a `MetricKind`") - return self._compiled.change_metric( + self._compiled.change_metric( metric_kind=metric_kind, metric_pointer=metric_pointer, metric_signature=metric_signature, @@ -1021,7 +1033,7 @@ def connectivity(self) -> int: :return: The connectivity of the index. :rtype: int """ - return self._compiled.connectivity + return int(self._compiled.connectivity) @property def capacity(self) -> int: @@ -1032,7 +1044,7 @@ def capacity(self) -> int: :return: The capacity of the index. :rtype: int """ - return self._compiled.capacity + return int(self._compiled.capacity) @property def memory_usage(self) -> int: @@ -1041,7 +1053,7 @@ def memory_usage(self) -> int: :return: The memory usage of the index. :rtype: int """ - return self._compiled.memory_usage + return int(self._compiled.memory_usage) @property def expansion_add(self) -> int: @@ -1052,10 +1064,10 @@ def expansion_add(self) -> int: :return: The expansion parameter for additions. :rtype: int """ - return self._compiled.expansion_add + return int(self._compiled.expansion_add) @expansion_add.setter - def expansion_add(self, v: int): + def expansion_add(self, v: int) -> None: """Sets the expansion parameter used during addition. :param v: The new expansion parameter for additions. @@ -1072,10 +1084,10 @@ def expansion_search(self) -> int: :return: The expansion parameter for searches. :rtype: int """ - return self._compiled.expansion_search + return int(self._compiled.expansion_search) @expansion_search.setter - def expansion_search(self, v: int): + def expansion_search(self, v: int) -> None: """Sets the expansion parameter used during searches. :param v: The new expansion parameter for searches. @@ -1085,7 +1097,7 @@ def expansion_search(self, v: int): def save( self, - path_or_buffer: str | os.PathLike | NoneType = None, + path_or_buffer: str | os.PathLike[str] | NoneType = None, progress: ProgressCallback | None = None, ) -> bytes | None: """Saves the index to a file or buffer. @@ -1103,7 +1115,7 @@ def save( path_or_buffer = path_or_buffer if path_or_buffer is not None else self.path if path_or_buffer is None: - return self._compiled.save_index_to_buffer(progress) + return cast("bytes", self._compiled.save_index_to_buffer(progress)) else: self._compiled.save_index_to_path(os.fspath(path_or_buffer), progress) return None @@ -1112,7 +1124,7 @@ def load( self, path_or_buffer: PathOrBuffer | NoneType = None, progress: ProgressCallback | None = None, - ): + ) -> None: """Loads the index from a file or buffer. If `path_or_buffer` is not provided, it defaults to the path stored in `self.path`. @@ -1132,17 +1144,17 @@ def load( if _is_buffer(path_or_buffer): self._compiled.load_index_from_buffer(path_or_buffer, progress) else: - path_or_buffer = os.fspath(path_or_buffer) - if os.path.exists(path_or_buffer): - self._compiled.load_index_from_path(path_or_buffer, progress) + path_str = os.fspath(cast("str | os.PathLike[str]", path_or_buffer)) + if os.path.exists(path_str): + self._compiled.load_index_from_path(path_str, progress) else: - raise FileNotFoundError(f"File not found: {path_or_buffer!r}") + raise FileNotFoundError(f"File not found: {path_str!r}") def view( self, path_or_buffer: PathOrBuffer | NoneType = None, progress: ProgressCallback | None = None, - ): + ) -> None: """Maps the index from a file or buffer without loading it into memory. If `path_or_buffer` is not provided, it defaults to the path stored in `self.path`. @@ -1161,19 +1173,19 @@ def view( if _is_buffer(path_or_buffer): self._compiled.view_index_from_buffer(path_or_buffer, progress) else: - self._compiled.view_index_from_path(os.fspath(path_or_buffer), progress) + self._compiled.view_index_from_path(os.fspath(cast("str | os.PathLike[str]", path_or_buffer)), progress) - def clear(self): + def clear(self) -> None: """Erases all vectors from the index, preserving the allocated space for future insertions.""" self._compiled.clear() - def reset(self): + def reset(self) -> None: """Erases all data from the index, closes any open files, and returns allocated memory to the OS.""" if not hasattr(self, "_compiled"): return self._compiled.reset() - def __del__(self): + def __del__(self) -> None: """Destructor method to reset the index when the object is deleted.""" self.reset() @@ -1220,11 +1232,14 @@ def join( """ assert not progress or _match_signature(progress, [int, int], bool), "Invalid callback signature" - return self._compiled.join( - other=other._compiled, - max_proposals=max_proposals, - exact=exact, - progress=progress, + return cast( + "dict[Key, Key]", + self._compiled.join( + other=other._compiled, + max_proposals=max_proposals, + exact=exact, + progress=progress, + ), ) def cluster( @@ -1304,11 +1319,11 @@ def pairwise_distance(self, left: KeyOrKeysLike, right: KeyOrKeysLike) -> NDArra assert isinstance(left, Iterable) == isinstance(right, Iterable) if not isinstance(left, Iterable): - return self._compiled.pairwise_distance(int(left), int(right)) # type: ignore[arg-type] + return cast("float", self._compiled.pairwise_distance(int(left), int(right))) # type: ignore[arg-type] else: left = np.array(left).astype(Key) right = np.array(right).astype(Key) - return self._compiled.pairwise_distances(left, right) + return cast("NDArray[Any]", self._compiled.pairwise_distances(left, right)) @property def keys(self) -> IndexedKeys: @@ -1335,7 +1350,7 @@ def max_level(self) -> int: :return: The maximum level in the graph. :rtype: int """ - return self._compiled.max_level + return int(self._compiled.max_level) @property def nlevels(self) -> int: @@ -1344,7 +1359,7 @@ def nlevels(self) -> int: :return: Number of levels in the graph. :rtype: int """ - return self._compiled.max_level + 1 + return int(self._compiled.max_level) + 1 @property def multi(self) -> bool: @@ -1353,7 +1368,7 @@ def multi(self) -> bool: :return: True if the index supports multi-value entries, False otherwise. :rtype: bool """ - return self._compiled.multi + return bool(self._compiled.multi) @property def stats(self) -> _CompiledIndexStats: @@ -1383,7 +1398,7 @@ def levels_stats(self) -> list[_CompiledIndexStats]: - `max_edges` (int): Maximum possible number of edges in the level. - `allocated_bytes` (int): Memory allocated for the level. """ - return self._compiled.levels_stats + return cast("list[Any]", self._compiled.levels_stats) def level_stats(self, level: int) -> _CompiledIndexStats: """Get statistics for a specific level of the graph. @@ -1486,7 +1501,7 @@ def __repr_pretty__(self) -> str: ) return lines - def _repr_pretty_(self, printer, cycle): + def _repr_pretty_(self, printer: Any, cycle: bool) -> None: """Handles pretty-printing of the object within interactive environments. :param printer: The pretty printer instance. @@ -1501,7 +1516,7 @@ class Indexes: def __init__( self, indexes: Iterable[Index] = [], - paths: Iterable[os.PathLike] = [], + paths: Iterable[os.PathLike[str]] = [], view: bool = False, threads: int = 0, ) -> None: @@ -1510,24 +1525,24 @@ def __init__( self._compiled.merge(index._compiled) self._compiled.merge_paths(paths, view=view, threads=threads) - def merge(self, index: Index): + def merge(self, index: Index) -> None: self._compiled.merge(index._compiled) - def merge_path(self, path: os.PathLike): + def merge_path(self, path: os.PathLike[str]) -> None: self._compiled.merge_path(os.fspath(path)) def __len__(self) -> int: - return self._compiled.__len__() + return int(self._compiled.__len__()) def search( self, - vectors, + vectors: VectorOrVectorsLike, count: int = 10, *, threads: int = 0, exact: bool = False, progress: ProgressCallback | None = None, - ): + ) -> Matches | BatchMatches: return _search_in_compiled( self._compiled.search_many, vectors, @@ -1616,19 +1631,19 @@ def search( else: raise ValueError("The `metric` must be a `CompiledMetric` or a `MetricKind`") - def search_batch(query, **kwargs): - assert dataset.shape[1] == query.shape[1], "Number of dimensions differs" - if dataset.dtype != query.dtype: - query = query.astype(dataset.dtype) + def search_batch(vectors: NDArray[Any], **kwargs: Any) -> _BatchSearchTuple: + assert dataset.shape[1] == vectors.shape[1], "Number of dimensions differs" + if dataset.dtype != vectors.dtype: + vectors = vectors.astype(dataset.dtype) - return _exact_search( + return cast(_BatchSearchTuple, _exact_search( dataset, - query, + vectors, metric_kind=metric_kind, metric_signature=metric_signature, metric_pointer=metric_pointer, **kwargs, - ) + )) return _search_in_compiled( search_batch, @@ -1643,8 +1658,8 @@ def search_batch(query, **kwargs): def kmeans( - X, - k, + X: NDArray[Any], + k: int, metric: str = "l2sq", dtype: str = "bf16", max_iterations: int = 300, @@ -1718,8 +1733,8 @@ def kmeans( >>> k = 5 >>> assignments, distances, centroids = usearch.index.kmeans(X, k) """ - metric = _normalize_metric(metric) - dtype = _normalize_dtype(dtype, ndim=X.shape[1], metric=metric) + metric_kind = cast("MetricKind", _normalize_metric(metric)) + dtype_kind = _normalize_dtype(dtype, ndim=X.shape[1], metric=metric_kind) # Generating a 64-bit unsigned integer in NumPy may be somewhat tricky. actual_seed: int | np.unsignedinteger = ( @@ -1728,12 +1743,12 @@ def kmeans( assignments, distances, centroids = _kmeans( X, k, - metric_kind=metric, + metric_kind=metric_kind, max_iterations=max_iterations, max_seconds=max_seconds, min_shifts=min_shifts, inertia_threshold=inertia_threshold, - dtype=dtype, + dtype=dtype_kind, seed=actual_seed, ) return assignments, distances, centroids diff --git a/python/usearch/io.py b/python/usearch/io.py index 8299d90d9..bb9e6c168 100644 --- a/python/usearch/io.py +++ b/python/usearch/io.py @@ -11,7 +11,7 @@ from numpy.typing import NDArray -def numpy_scalar_size(dtype) -> int: +def numpy_scalar_size(dtype: Any) -> int: return { np.float64: 8, np.int64: 8, @@ -113,7 +113,7 @@ def load_matrix( ).reshape(rows, cols) -def save_matrix(vectors: NDArray[Any], filename: str): +def save_matrix(vectors: NDArray[Any], filename: str) -> None: """Write *.ibin, *.bbib, *.hbin, *.fbin, *.dbin, *.i8bin, *.i32bin, *.f32bin files with matrices. :param vectors: the matrix to serialize diff --git a/python/usearch/numba.py b/python/usearch/numba.py index 40d53a52a..3a82352ce 100644 --- a/python/usearch/numba.py +++ b/python/usearch/numba.py @@ -3,6 +3,7 @@ # Python tooling, linters, and static analyzers. It also embeds JIT # into the primary `Index` class, connecting USearch with Numba. from math import sqrt +from typing import Any from usearch.index import CompiledMetric, MetricKind, MetricSignature, ScalarKind @@ -11,7 +12,7 @@ def jit( ndim: int, metric: MetricKind = MetricKind.Cos, dtype: ScalarKind = ScalarKind.F32, -) -> CompiledMetric: +) -> CompiledMetric | MetricKind: """JIT-compiles the metric for target hardware and number of dimensions. This can result in up-to 3x performance difference on very large vectors @@ -51,7 +52,7 @@ def jit( } accumulator = scalar_kind_to_accumulator_type[dtype] - def numba_ip(a, b): + def numba_ip(a: Any, b: Any) -> Any: a_array = carray(a, ndim) b_array = carray(b, ndim) ab = accumulator(0) @@ -59,7 +60,7 @@ def numba_ip(a, b): ab += a_array[i] * b_array[i] return types.float32(1 - ab) - def numba_cos(a, b): + def numba_cos(a: Any, b: Any) -> Any: a_array = carray(a, ndim) b_array = carray(b, ndim) ab = accumulator(0) @@ -78,7 +79,7 @@ def numba_cos(a, b): else: return types.float32(1 - ab / (a_norm * b_norm)) - def numba_l2sq(a, b): + def numba_l2sq(a: Any, b: Any) -> Any: a_array = carray(a, ndim) b_array = carray(b, ndim) ab_delta_sq = accumulator(0) diff --git a/python/usearch/server.py b/python/usearch/server.py index d30e94a1e..363213b63 100644 --- a/python/usearch/server.py +++ b/python/usearch/server.py @@ -3,10 +3,10 @@ import argparse import os -from typing import TYPE_CHECKING, Any +from typing import TYPE_CHECKING, Any, cast import numpy as np -from ucall.rich_posix import Server # type: ignore[import-untyped] +from ucall.rich_posix import Server # type: ignore[import-not-found] from usearch.index import Index, Key @@ -38,7 +38,7 @@ def serve( threads: int = 1, path: str = "index.usearch", immutable: bool = False, -): +) -> None: server = Server(port=port) index = Index(ndim=ndim_, metric=metric) @@ -48,50 +48,50 @@ def serve( else: index.load(path) - @server + @server # type: ignore[untyped-decorator] def size() -> int: return len(index) - @server + @server # type: ignore[untyped-decorator] def ndim() -> int: return index.ndim - @server + @server # type: ignore[untyped-decorator] def capacity() -> int: return index.capacity - @server + @server # type: ignore[untyped-decorator] def connectivity() -> int: return index.connectivity - @server - def add_one(key: int, vector: NDArray[Any]): + @server # type: ignore[untyped-decorator] + def add_one(key: int, vector: NDArray[Any]) -> None: keys = np.array([key], dtype=Key) vectors = vector.flatten().reshape(vector.shape[0], 1) index.add(keys, vectors) - @server - def add_many(keys: NDArray[Any], vectors: NDArray[Any]): + @server # type: ignore[untyped-decorator] + def add_many(keys: NDArray[Any], vectors: NDArray[Any]) -> None: index.add(keys, vectors, threads=threads) - @server - def search_one(vector: NDArray[Any], count: int) -> list[tuple]: + @server # type: ignore[untyped-decorator] + def search_one(vector: NDArray[Any], count: int) -> list[tuple[int, float]]: vectors = vector.reshape(vector.shape[0], 1) results: Matches | BatchMatches = index.search(vectors, count) return results.to_list() - @server - def search_many(vectors: NDArray[Any], count: int) -> list[tuple]: + @server # type: ignore[untyped-decorator] + def search_many(vectors: NDArray[Any], count: int) -> list[tuple[int, float]]: results: Matches | BatchMatches = index.search(vectors, count) return results.to_list() - @server - def add_ascii(key: int, string: str): - return add_one(key, _ascii_to_vector(string)) + @server # type: ignore[untyped-decorator] + def add_ascii(key: int, string: str) -> None: + add_one(key, _ascii_to_vector(string)) - @server - def search_ascii(string: str, count: int): - return search_one(_ascii_to_vector(string), count) + @server # type: ignore[untyped-decorator] + def search_ascii(string: str, count: int) -> list[tuple[int, float]]: + return cast("list[tuple[int, float]]", search_one(_ascii_to_vector(string), count)) try: server.run()