ENH: Add annotations for scipy.spatial.cKDTree

mypy.ini

@@ -103,9 +103,6 @@ ignore_missing_imports = True
 [mypy-scipy.interpolate.interpnd]
 ignore_missing_imports = True
 
-[mypy-scipy.spatial.ckdtree]
-ignore_missing_imports = True
-
 [mypy-scipy.linalg.cython_blas]
 ignore_missing_imports = True
 

scipy/spatial/ckdtree.pyi

@@ -1,9 +1,221 @@
+import sys
+from typing import (
+    Any,
+    Dict,
+    Generic,
+    List,
+    Optional,
+    overload,
+    Set,
+    Tuple,
+    TypeVar,
+    Union,
+)
+
 import numpy as np
+import numpy.typing as npt
+from scipy.sparse import coo_matrix, dok_matrix
+
+if sys.version_info >= (3, 8):
+    from typing import Literal
+else:
+    from typing_extensions import Literal
+
+# TODO: Replace `ndarray` with a 1D float64 array when possible
+_BoxType = TypeVar("_BoxType", None, np.ndarray)
+
+# Copied from `numpy.typing._scalar_like._ScalarLike`
+# TODO: Expand with 0D arrays once we have shape support
+_ArrayLike0D = Union[
+    bool,
+    int,
+    float,
+    complex,
+    str,
+    bytes,
+    np.generic,
+]
+_WeightType = Union[npt.ArrayLike, Tuple[npt.ArrayLike, npt.ArrayLike]]
 
 class cKDTreeNode:
     @property
     def data_points(self) -> np.ndarray: ...
     @property
     def indices(self) -> np.ndarray: ...
 
-class cKDTree: ...
+    # These are read-only attributes in cython, which behave like properties
+    @property
+    def level(self) -> int: ...
+    @property
+    def split_dim(self) -> int: ...
+    @property
+    def children(self) -> int: ...
+    @property
+    def start_idx(self) -> int: ...
+    @property
+    def end_idx(self) -> int: ...
+    @property
+    def split(self) -> float: ...
+    @property
+    def lesser(self) -> Optional[cKDTreeNode]: ...
+    @property
+    def greater(self) -> Optional[cKDTreeNode]: ...
+
+class cKDTree(Generic[_BoxType]):
+    @property
+    def n(self) -> int: ...
+    @property
+    def m(self) -> int: ...
+    @property
+    def leafsize(self) -> int: ...
+    @property
+    def size(self) -> int: ...
+    @property
+    def tree(self) -> cKDTreeNode: ...
+
+    # These are read-only attributes in cython, which behave like properties
+    @property
+    def data(self) -> np.ndarray: ...
+    @property
+    def maxes(self) -> np.ndarray: ...
+    @property
+    def mins(self) -> np.ndarray: ...
+    @property
+    def indices(self) -> np.ndarray: ...
+    @property
+    def boxsize(self) -> _BoxType: ...
+
+    # NOTE: In practice `__init__` is used as constructor, not `__new__`.
+    # The latter gives us more flexibility in setting the generic parameter
+    # though.
+    @overload
+    def __new__(  # type: ignore[misc]
+        cls,
+        data: npt.ArrayLike,
+        leafsize: int = ...,
+        compact_nodes: bool = ...,
+        copy_data: bool = ...,
+        balanced_tree: bool = ...,
+        boxsize: None = ...,
+    ) -> cKDTree[None]: ...
+    @overload
+    def __new__(
+        cls,
+        data: npt.ArrayLike,
+        leafsize: int = ...,
+        compact_nodes: bool = ...,
+        copy_data: bool = ...,
+        balanced_tree: bool = ...,
+        boxsize: npt.ArrayLike = ...,
+    ) -> cKDTree[np.ndarray]: ...
+
+    # TODO: returns a 2-tuple of scalars if `x.ndim == 1` and `k == 1`,
+    # returns a 2-tuple of arrays otherwise
+    def query(
+        self,
+        x: npt.ArrayLike,
+        k: npt.ArrayLike = ...,
+        eps: float = ...,
+        p: float = ...,
+        distance_upper_bound: float = ...,
+        workers: Optional[int] = ...,
+    ) -> Tuple[Any, Any]: ...
+
+    # TODO: returns a list scalars if `x.ndim <= 1`,
+    # returns an object array of lists otherwise
+    def query_ball_point(
+        self,
+        x: npt.ArrayLike,
+        r: npt.ArrayLike,
+        p: float,
+        eps: float = ...,
+        workers: Optional[int] = ...,
+        return_sorted: Optional[bool] = ...,
+        return_length: bool = ...
+    ) -> Any: ...
+
+    def query_ball_tree(
+        self,
+        other: cKDTree,
+        r: float,
+        p: float,
+        eps: float = ...,
+    ) -> List[List[int]]: ...
+
+    @overload
+    def query_pairs(  # type: ignore[misc]
+        self,
+        r: float,
+        p: float = ...,
+        eps: float = ...,
+        output_type: Literal["set"] = ...,
+    ) -> Set[Tuple[int, int]]: ...
+    @overload
+    def query_pairs(
+        self,
+        r: float,
+        p: float = ...,
+        eps: float = ...,
+        output_type: Literal["ndarray"] = ...,
+    ) -> np.ndarray: ...
+
+    @overload
+    def count_neighbors(  # type: ignore[misc]
+        self,
+        other: cKDTree,
+        r: _ArrayLike0D,
+        p: float = ...,
+        weights: None = ...,
+        cumulative: bool = ...,
+    ) -> int: ...
+    @overload
+    def count_neighbors(  # type: ignore[misc]
+        self,
+        other: cKDTree,
+        r: _ArrayLike0D,
+        p: float = ...,
+        weights: _WeightType = ...,
+        cumulative: bool = ...,
+    ) -> np.float64: ...
+    @overload
+    def count_neighbors(
+        self,
+        other: cKDTree,
+        r: npt.ArrayLike,
+        p: float = ...,
+        weights: Optional[_WeightType] = ...,
+        cumulative: bool = ...,
+    ) -> np.ndarray: ...
+
+    @overload
+    def sparse_distance_matrix(  # type: ignore[misc]
+        self,
+        other: cKDTree,
+        max_distance: float,
+        p: float = ...,
+        output_type: Literal["dok_matrix"] = ...,
+    ) -> dok_matrix: ...
+    @overload
+    def sparse_distance_matrix(  # type: ignore[misc]
+        self,
+        other: cKDTree,
+        max_distance: float,
+        p: float = ...,
+        output_type: Literal["coo_matrix"] = ...,
+    ) -> coo_matrix: ...
+    @overload
+    def sparse_distance_matrix(  # type: ignore[misc]
+        self,
+        other: cKDTree,
+        max_distance: float,
+        p: float = ...,
+        output_type: Literal["dict"] = ...,
+    ) -> Dict[Tuple[int, int], float]: ...
+    @overload
+    def sparse_distance_matrix(
+        self,
+        other: cKDTree,
+        max_distance: float,
+        p: float = ...,
+        output_type: Literal["ndarray"] = ...,
+    ) -> np.ndarray: ...