Make type annotations for NumPy arrays more specific (#1358)

* Make type annotations for NumPy arrays more specific

* Add np.generic to nd.dtype annotations

* Use Any instead of np.generic

Author: IvanIsCoding

rustworkx/__init__.pyi

@@ -11,6 +11,7 @@
 
 import sys
 import numpy as np
+import numpy.typing as npt
 
 from typing import Generic, Any, Callable, overload
 from collections.abc import Iterator, Sequence
@@ -289,7 +290,7 @@ def distance_matrix(
     parallel_threshold: int = ...,
     as_undirected: bool = ...,
     null_value: float = ...,
-) -> np.ndarray: ...
+) -> npt.NDArray[np.float64]: ...
 def unweighted_average_shortest_path_length(
     graph: PyGraph | PyDiGraph,
     parallel_threshold: int = ...,
@@ -300,7 +301,7 @@ def adjacency_matrix(
     weight_fn: Callable[[_T], float] | None = ...,
     default_weight: float = ...,
     null_value: float = ...,
-) -> np.ndarray: ...
+) -> npt.NDArray[np.float64]: ...
 def all_simple_paths(
     graph: PyGraph | PyDiGraph,
     from_: int,
@@ -319,13 +320,13 @@ def floyd_warshall_numpy(
     weight_fn: Callable[[_T], float] | None = ...,
     default_weight: float = ...,
     parallel_threshold: int = ...,
-) -> np.ndarray: ...
+) -> npt.NDArray[np.float64]: ...
 def floyd_warshall_successor_and_distance(
     graph: PyGraph[_S, _T] | PyDiGraph[_S, _T],
     weight_fn: Callable[[_T], float] | None = ...,
     default_weight: float | None = ...,
     parallel_threshold: int | None = ...,
-) -> tuple[np.ndarray, np.ndarray]: ...
+) -> tuple[npt.NDArray[np.float64], npt.NDArray[np.float64]]: ...
 def astar_shortest_path(
     graph: PyGraph[_S, _T] | PyDiGraph[_S, _T],
     node: int,

rustworkx/rustworkx.pyi

@@ -35,6 +35,7 @@ from rustworkx import generators  # noqa
 from typing_extensions import Self
 
 import numpy as np
+import numpy.typing as npt
 import sys
 
 if sys.version_info >= (3, 13):
@@ -206,15 +207,15 @@ def digraph_adjacency_matrix(
     default_weight: float = ...,
     null_value: float = ...,
     parallel_edge: str = ...,
-) -> np.ndarray: ...
+) -> npt.NDArray[np.float64]: ...
 def graph_adjacency_matrix(
     graph: PyGraph[_S, _T],
     /,
     weight_fn: Callable[[_T], float] | None = ...,
     default_weight: float = ...,
     null_value: float = ...,
     parallel_edge: str = ...,
-) -> np.ndarray: ...
+) -> npt.NDArray[np.float64]: ...
 def cycle_basis(graph: PyGraph, /, root: int | None = ...) -> list[list[int]]: ...
 def articulation_points(graph: PyGraph, /) -> set[int]: ...
 def bridges(graph: PyGraph, /) -> set[tuple[int]]: ...
@@ -595,14 +596,14 @@ def undirected_gnp_random_graph(
 ) -> PyGraph: ...
 def directed_sbm_random_graph(
     sizes: list[int],
-    probabilities: np.ndarray,
+    probabilities: npt.NDArray[np.float64],
     loops: bool,
     /,
     seed: int | None = ...,
 ) -> PyDiGraph: ...
 def undirected_sbm_random_graph(
     sizes: list[int],
-    probabilities: np.ndarray,
+    probabilities: npt.NDArray[np.float64],
     loops: bool,
     /,
     seed: int | None = ...,
@@ -863,13 +864,13 @@ def digraph_distance_matrix(
     parallel_threshold: int | None = ...,
     as_undirected: bool | None = ...,
     null_value: float | None = ...,
-) -> np.ndarray: ...
+) -> npt.NDArray[np.float64]: ...
 def graph_distance_matrix(
     graph: PyGraph,
     /,
     parallel_threshold: int | None = ...,
     null_value: float | None = ...,
-) -> np.ndarray: ...
+) -> npt.NDArray[np.float64]: ...
 def digraph_floyd_warshall(
     graph: PyDiGraph[_S, _T],
     /,
@@ -892,29 +893,29 @@ def digraph_floyd_warshall_numpy(
     as_undirected: bool | None = ...,
     default_weight: float | None = ...,
     parallel_threshold: int | None = ...,
-) -> np.ndarray: ...
+) -> npt.NDArray[np.float64]: ...
 def graph_floyd_warshall_numpy(
     graph: PyGraph[_S, _T],
     /,
     weight_fn: Callable[[_T], float] | None = ...,
     default_weight: float | None = ...,
     parallel_threshold: int | None = ...,
-) -> np.ndarray: ...
+) -> npt.NDArray[np.float64]: ...
 def digraph_floyd_warshall_successor_and_distance(
     graph: PyDiGraph[_S, _T],
     /,
     weight_fn: Callable[[_T], float] | None = ...,
     as_undirected: bool | None = ...,
     default_weight: float | None = ...,
     parallel_threshold: int | None = ...,
-) -> tuple[np.ndarray, np.ndarray]: ...
+) -> tuple[npt.NDArray[np.float64], npt.NDArray[np.float64]]: ...
 def graph_floyd_warshall_successor_and_distance(
     graph: PyGraph[_S, _T],
     /,
     weight_fn: Callable[[_T], float] | None = ...,
     default_weight: float | None = ...,
     parallel_threshold: int | None = ...,
-) -> tuple[np.ndarray, np.ndarray]: ...
+) -> tuple[npt.NDArray[np.float64], npt.NDArray[np.float64]]: ...
 def find_negative_cycle(
     graph: PyDiGraph[_S, _T],
     edge_cost_fn: Callable[[_T], float],
@@ -1079,7 +1080,9 @@ class _RustworkxCustomVecIter(Generic[_T_co], Sequence[_T_co], ABC):
     def __len__(self) -> int: ...
     def __ne__(self, other: object) -> bool: ...
     def __setstate__(self, state: Sequence[_T_co]) -> None: ...
-    def __array__(self, dtype: np.dtype | None = ..., copy: bool | None = ...) -> np.ndarray: ...
+    def __array__(
+        self, dtype: np.dtype[Any] | None = ..., copy: bool | None = ...
+    ) -> npt.NDArray[Any]: ...
     def __iter__(self) -> Iterator[_T_co]: ...
     def __reversed__(self) -> Iterator[_T_co]: ...
 
@@ -1235,11 +1238,11 @@ class PyGraph(Generic[_S, _T]):
     ) -> int | None: ...
     @staticmethod
     def from_adjacency_matrix(
-        matrix: np.ndarray, /, null_value: float = ...
+        matrix: npt.NDArray[np.float64], /, null_value: float = ...
     ) -> PyGraph[int, float]: ...
     @staticmethod
     def from_complex_adjacency_matrix(
-        matrix: np.ndarray, /, null_value: complex = ...
+        matrix: npt.NDArray[np.complex64], /, null_value: complex = ...
     ) -> PyGraph[int, complex]: ...
     def get_all_edge_data(self, node_a: int, node_b: int, /) -> list[_T]: ...
     def get_edge_data(self, node_a: int, node_b: int, /) -> _T: ...
@@ -1400,11 +1403,11 @@ class PyDiGraph(Generic[_S, _T]):
     ) -> list[_S]: ...
     @staticmethod
     def from_adjacency_matrix(
-        matrix: np.ndarray, /, null_value: float = ...
+        matrix: npt.NDArray[np.float64], /, null_value: float = ...
     ) -> PyDiGraph[int, float]: ...
     @staticmethod
     def from_complex_adjacency_matrix(
-        matrix: np.ndarray, /, null_value: complex = ...
+        matrix: npt.NDArray[np.complex64], /, null_value: complex = ...
     ) -> PyDiGraph[int, complex]: ...
     def get_all_edge_data(self, node_a: int, node_b: int, /) -> list[_T]: ...
     def get_edge_data(self, node_a: int, node_b: int, /) -> _T: ...