Merge pull request #74 from jorenham/stubtest/scipy.spatial

`scipy.spatial`

.github/workflows/ci.yml

@@ -57,6 +57,7 @@ jobs:
             "misc"
             "ndimage"
             "odr"
+            "spatial"
             "special"
           )
           for target in "${targets[@]}"; do
@@ -88,7 +89,7 @@ jobs:
           "scipy.optimize",
           # "scipy.signal",
           # "scipy.sparse",
-          # "scipy.spatial",
+          "scipy.spatial",
           "scipy.special",
           # "scipy.stats",
         ]

README.md

@@ -119,6 +119,6 @@ The exact version requirements are specified in the [`pyproject.toml`](pyproject
 | `scipy.optimize`    | :heavy_check_mark:  | :heavy_check_mark: | :x:                   | :first_quarter_moon:   |
 | `scipy.signal`      | :heavy_check_mark:  | :x:                | :x:                   | :new_moon:             |
 | `scipy.sparse`      | :heavy_check_mark:  | :x:                | :x:                   | :waxing_crescent_moon: |
-| `scipy.spatial`     | :heavy_check_mark:  | :x:                | :x:                   | :waxing_crescent_moon: |
+| `scipy.spatial`     | :heavy_check_mark:  | :heavy_check_mark: | :heavy_check_mark:    | :waxing_gibbous_moon:  |
 | `scipy.special`     | :heavy_check_mark:  | :heavy_check_mark: | :heavy_check_mark:    | :first_quarter_moon:   |
 | `scipy.stats`       | :heavy_check_mark:  | :x:                | :x:                   | :first_quarter_moon:   |

scipy-stubs/spatial/__init__.pyi

@@ -1,12 +1,35 @@
-from scipy._lib._testutils import PytestTester as PytestTester
-from scipy._typing import Untyped
-from . import ckdtree as ckdtree, distance as distance, kdtree as kdtree, qhull as qhull, transform as transform
-from ._ckdtree import *
-from ._geometric_slerp import geometric_slerp as geometric_slerp
-from ._kdtree import *
-from ._plotutils import *
-from ._procrustes import procrustes as procrustes
-from ._qhull import *
-from ._spherical_voronoi import SphericalVoronoi as SphericalVoronoi
+from . import ckdtree, distance, kdtree, qhull, transform
+from ._ckdtree import cKDTree
+from ._geometric_slerp import geometric_slerp
+from ._kdtree import KDTree, Rectangle, distance_matrix, minkowski_distance, minkowski_distance_p
+from ._plotutils import convex_hull_plot_2d, delaunay_plot_2d, voronoi_plot_2d
+from ._procrustes import procrustes
+from ._qhull import ConvexHull, Delaunay, HalfspaceIntersection, QhullError, Voronoi, tsearch
+from ._spherical_voronoi import SphericalVoronoi
 
-test: Untyped
+__all__ = [
+    "ConvexHull",
+    "Delaunay",
+    "HalfspaceIntersection",
+    "KDTree",
+    "QhullError",
+    "Rectangle",
+    "SphericalVoronoi",
+    "Voronoi",
+    "cKDTree",
+    "ckdtree",
+    "convex_hull_plot_2d",
+    "delaunay_plot_2d",
+    "distance",
+    "distance",
+    "distance_matrix",
+    "geometric_slerp",
+    "kdtree",
+    "minkowski_distance",
+    "minkowski_distance_p",
+    "procrustes",
+    "qhull",
+    "transform",
+    "tsearch",
+    "voronoi_plot_2d",
+]

scipy-stubs/spatial/_ckdtree.pyi

@@ -1,18 +1,24 @@
-# NOTE: There are alteady inline stubs for this module, so this won't be used (which is unfortunate, because the inline
-# annotations are wrong in several places, although it's still pretty good overall).
-from typing import Any, Generic, Literal, TypeAlias, TypeVar, overload
+# NOTE(jorenham): There are alteady inline stubs for this module, so this won't be used (which is unfortunate, because the inline
+# annotations are wrong in several places, although I've seen worse).
+
+from typing import Literal, TypeAlias, overload, type_check_only
 
 import numpy as np
 import numpy.typing as npt
+from numpy._typing import _ArrayLikeInt
 from scipy._typing import AnyScalar
 from scipy.sparse import coo_matrix, dok_matrix
 
 __all__ = ["cKDTree"]
 
 _Weights: TypeAlias = npt.ArrayLike | tuple[npt.ArrayLike, npt.ArrayLike]
-_BoxT = TypeVar("_BoxT", bound=npt.NDArray[np.float64] | None)
 
-class cKDTreeNode:
+@type_check_only
+class _CythonMixin:
+    def __setstate_cython__(self, pyx_state: object, /) -> None: ...
+    def __reduce_cython__(self, /) -> None: ...
+
+class cKDTreeNode(_CythonMixin):
     @property
     def data_points(self) -> npt.NDArray[np.float64]: ...
     @property
@@ -36,7 +42,7 @@ class cKDTreeNode:
     @property
     def greater(self) -> cKDTreeNode | None: ...
 
-class cKDTree(Generic[_BoxT]):
+class cKDTree(_CythonMixin):
     @property
     def n(self) -> int: ...
     @property
@@ -58,147 +64,158 @@ class cKDTree(Generic[_BoxT]):
     @property
     def indices(self) -> npt.NDArray[np.float64]: ...
     @property
-    def boxsize(self) -> _BoxT: ...
+    def boxsize(self) -> npt.NDArray[np.float64] | None: ...
 
-    # 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__(
-        cls,
-        data: npt.ArrayLike,
-        leafsize: int = ...,
-        compact_nodes: bool = ...,
-        copy_data: bool = ...,
-        balanced_tree: bool = ...,
-        boxsize: None = ...,
-    ) -> cKDTree[None]: ...
-    @overload
-    def __new__(
-        cls,
+    #
+    def __init__(
+        self,
+        /,
         data: npt.ArrayLike,
         leafsize: int = ...,
         compact_nodes: bool = ...,
         copy_data: bool = ...,
         balanced_tree: bool = ...,
-        boxsize: npt.ArrayLike = ...,
-    ) -> cKDTree[npt.NDArray[np.float64]]: ...
-
-    # TODO: returns a 2-tuple of scalars if `x.ndim == 1` and `k == 1`,
-    # returns a 2-tuple of arrays otherwise
+        boxsize: npt.ArrayLike | None = ...,
+    ) -> None: ...
     def query(
         self,
         x: npt.ArrayLike,
-        k: npt.ArrayLike = ...,
-        eps: float = ...,
-        p: float = ...,
-        distance_upper_bound: float = ...,
-        workers: int | None = ...,
-    ) -> tuple[Any, Any]: ...
-
-    # TODO: returns a list scalars if `x.ndim <= 1`,
-    # returns an object array of lists otherwise
+        k: _ArrayLikeInt = 1,
+        eps: float = 0.0,
+        p: float = 2.0,
+        distance_upper_bound: float = ...,  # inf
+        workers: int | None = None,
+    ) -> tuple[float, float] | tuple[npt.NDArray[np.float64], npt.NDArray[np.float64]]: ...
     def query_ball_point(
         self,
         x: npt.ArrayLike,
         r: npt.ArrayLike,
-        p: float = ...,
-        eps: float = ...,
-        workers: int | None = ...,
-        return_sorted: bool | None = ...,
-        return_length: bool = ...,
-    ) -> Any: ...
+        p: float = 2.0,
+        eps: float = 0.0,
+        workers: int | None = None,
+        return_sorted: bool | None = None,
+        return_length: bool = False,
+    ) -> list[int] | npt.NDArray[np.object_]: ...
     def query_ball_tree(
         self,
-        other: cKDTree[npt.NDArray[np.float64] | None],
+        other: cKDTree,
         r: float,
-        p: float,
-        eps: float = ...,
+        p: float = 2.0,
+        eps: float = 0.0,
     ) -> list[list[int]]: ...
     @overload
     def query_pairs(
         self,
         r: float,
-        p: float = ...,
-        eps: float = ...,
-        output_type: Literal["set"] = ...,
+        p: float = 2.0,
+        eps: float = 0,
+        output_type: Literal["set"] = "set",
     ) -> set[tuple[int, int]]: ...
     @overload
     def query_pairs(
         self,
         r: float,
-        p: float = ...,
-        eps: float = ...,
-        output_type: Literal["ndarray"] = ...,
+        p: float = 2.0,
+        eps: float = 0,
+        *,
+        output_type: Literal["ndarray"],
+    ) -> npt.NDArray[np.intp]: ...
+    @overload
+    def query_pairs(
+        self,
+        r: float,
+        p: float,
+        eps: float,
+        output_type: Literal["ndarray"],
     ) -> npt.NDArray[np.intp]: ...
     @overload
     def count_neighbors(
         self,
-        other: cKDTree[npt.NDArray[np.float64] | None],
+        other: cKDTree,
         r: AnyScalar,
-        p: float = ...,
-        weights: None | tuple[None, None] = ...,
-        cumulative: bool = ...,
+        p: float = 2.0,
+        weights: tuple[None, None] | None = None,
+        cumulative: bool = True,
     ) -> np.intp: ...
     @overload
     def count_neighbors(
         self,
-        other: cKDTree[npt.NDArray[np.float64] | None],
+        other: cKDTree,
+        r: AnyScalar,
+        p: float,
+        weights: _Weights,
+        cumulative: bool = True,
+    ) -> np.float64: ...
+    @overload
+    def count_neighbors(
+        self,
+        other: cKDTree,
         r: AnyScalar,
-        p: float = ...,
-        weights: _Weights = ...,
-        cumulative: bool = ...,
+        p: float = 2.0,
+        *,
+        weights: _Weights,
+        cumulative: bool = True,
     ) -> np.float64: ...
     @overload
     def count_neighbors(
         self,
-        other: cKDTree[npt.NDArray[np.float64] | None],
+        other: cKDTree,
         r: npt.ArrayLike,
-        p: float = ...,
+        p: float = 2.0,
         weights: None | tuple[None, None] = ...,
-        cumulative: bool = ...,
+        cumulative: bool = True,
     ) -> np.float64 | np.intp | npt.NDArray[np.intp]: ...
     @overload
     def count_neighbors(
         self,
-        other: cKDTree[npt.NDArray[np.float64] | None],
+        other: cKDTree,
         r: npt.ArrayLike,
-        p: float = ...,
-        weights: _Weights = ...,
-        cumulative: bool = ...,
-    ) -> np.intp | np.float64 | npt.NDArray[np.float64]: ...
+        p: float,
+        weights: _Weights,
+        cumulative: bool = True,
+    ) -> np.float64 | np.intp | npt.NDArray[np.float64]: ...
+    @overload
+    def count_neighbors(
+        self,
+        other: cKDTree,
+        r: npt.ArrayLike,
+        p: float = 2.0,
+        *,
+        weights: _Weights,
+        cumulative: bool = True,
+    ) -> np.float64 | np.intp | npt.NDArray[np.float64]: ...
     @overload
     def sparse_distance_matrix(
         self,
-        other: cKDTree[npt.NDArray[np.float64] | None],
+        other: cKDTree,
         max_distance: float,
-        p: float = ...,
+        p: float = 2.0,
         output_type: Literal["dok_matrix"] = ...,
     ) -> dok_matrix: ...
     @overload
     def sparse_distance_matrix(
         self,
-        other: cKDTree[npt.NDArray[np.float64] | None],
+        other: cKDTree,
         max_distance: float,
-        p: float = ...,
+        p: float = 2.0,
         *,
         output_type: Literal["coo_matrix"],
     ) -> coo_matrix: ...
     @overload
     def sparse_distance_matrix(
         self,
-        other: cKDTree[npt.NDArray[np.float64] | None],
+        other: cKDTree,
         max_distance: float,
-        p: float = ...,
+        p: float = 2.0,
         *,
         output_type: Literal["dict"],
     ) -> dict[tuple[int, int], float]: ...
     @overload
     def sparse_distance_matrix(
         self,
-        other: cKDTree[npt.NDArray[np.float64] | None],
+        other: cKDTree,
         max_distance: float,
-        p: float = ...,
+        p: float = 2.0,
         *,
         output_type: Literal["ndarray"],
     ) -> npt.NDArray[np.void]: ...