Merge pull request #19683 from BvB93/stride_tricks

ENH: Add annotations for `np.lib.stride_tricks`

numpy/lib/stride_tricks.pyi

@@ -1,16 +1,82 @@
-from typing import Any, List
+from typing import Any, List, Dict, Iterable, TypeVar, overload
+from typing_extensions import SupportsIndex
 
-from numpy.typing import _ShapeLike, _Shape
+from numpy import dtype, generic
+from numpy.typing import (
+    NDArray,
+    ArrayLike,
+    _ShapeLike,
+    _Shape,
+    _NestedSequence,
+    _SupportsArray,
+)
+
+_SCT = TypeVar("_SCT", bound=generic)
+_ArrayLike = _NestedSequence[_SupportsArray[dtype[_SCT]]]
 
 __all__: List[str]
 
 class DummyArray:
-    __array_interface__: Any
-    base: Any
-    def __init__(self, interface, base=...): ...
+    __array_interface__: Dict[str, Any]
+    base: None | NDArray[Any]
+    def __init__(
+        self,
+        interface: Dict[str, Any],
+        base: None | NDArray[Any] = ...,
+    ) -> None: ...
+
+@overload
+def as_strided(
+    x: _ArrayLike[_SCT],
+    shape: None | Iterable[int] = ...,
+    strides: None | Iterable[int] = ...,
+    subok: bool = ...,
+    writeable: bool = ...,
+) -> NDArray[_SCT]: ...
+@overload
+def as_strided(
+    x: ArrayLike,
+    shape: None | Iterable[int] = ...,
+    strides: None | Iterable[int] = ...,
+    subok: bool = ...,
+    writeable: bool = ...,
+) -> NDArray[Any]: ...
+
+@overload
+def sliding_window_view(
+    x: _ArrayLike[_SCT],
+    window_shape: int | Iterable[int],
+    axis: None | SupportsIndex = ...,
+    *,
+    subok: bool = ...,
+    writeable: bool = ...,
+) -> NDArray[_SCT]: ...
+@overload
+def sliding_window_view(
+    x: ArrayLike,
+    window_shape: int | Iterable[int],
+    axis: None | SupportsIndex = ...,
+    *,
+    subok: bool = ...,
+    writeable: bool = ...,
+) -> NDArray[Any]: ...
+
+@overload
+def broadcast_to(
+    array: _ArrayLike[_SCT],
+    shape: int | Iterable[int],
+    subok: bool = ...,
+) -> NDArray[_SCT]: ...
+@overload
+def broadcast_to(
+    array: ArrayLike,
+    shape: int | Iterable[int],
+    subok: bool = ...,
+) -> NDArray[Any]: ...
 
-def as_strided(x, shape=..., strides=..., subok=..., writeable=...): ...
-def sliding_window_view(x, window_shape, axis=..., *, subok=..., writeable=...): ...
-def broadcast_to(array, shape, subok=...): ...
 def broadcast_shapes(*args: _ShapeLike) -> _Shape: ...
-def broadcast_arrays(*args, subok=...): ...
+
+def broadcast_arrays(
+    *args: ArrayLike,
+    subok: bool = ...,
+) -> List[NDArray[Any]]: ...

numpy/typing/tests/data/fail/stride_tricks.py

@@ -0,0 +1,9 @@
+import numpy as np
+import numpy.typing as npt
+
+AR_f8: npt.NDArray[np.float64]
+
+np.lib.stride_tricks.as_strided(AR_f8, shape=8)  # E: No overload variant
+np.lib.stride_tricks.as_strided(AR_f8, strides=8)  # E: No overload variant
+
+np.lib.stride_tricks.sliding_window_view(AR_f8, axis=(1,))  # E: No overload variant

numpy/typing/tests/data/reveal/stride_tricks.py

@@ -0,0 +1,28 @@
+from typing import List, Dict, Any
+import numpy as np
+import numpy.typing as npt
+
+AR_f8: npt.NDArray[np.float64]
+AR_LIKE_f: List[float]
+interface_dict: Dict[str, Any]
+
+reveal_type(np.lib.stride_tricks.DummyArray(interface_dict))  # E: numpy.lib.stride_tricks.DummyArray
+
+reveal_type(np.lib.stride_tricks.as_strided(AR_f8))  # E: numpy.ndarray[Any, numpy.dtype[{float64}]]
+reveal_type(np.lib.stride_tricks.as_strided(AR_LIKE_f))  # E: numpy.ndarray[Any, numpy.dtype[Any]]
+reveal_type(np.lib.stride_tricks.as_strided(AR_f8, strides=(1, 5)))  # E: numpy.ndarray[Any, numpy.dtype[{float64}]]
+reveal_type(np.lib.stride_tricks.as_strided(AR_f8, shape=[9, 20]))  # E: numpy.ndarray[Any, numpy.dtype[{float64}]]
+
+reveal_type(np.lib.stride_tricks.sliding_window_view(AR_f8, 5))  # E: numpy.ndarray[Any, numpy.dtype[{float64}]]
+reveal_type(np.lib.stride_tricks.sliding_window_view(AR_LIKE_f, (1, 5)))  # E: numpy.ndarray[Any, numpy.dtype[Any]]
+reveal_type(np.lib.stride_tricks.sliding_window_view(AR_f8, [9], axis=1))  # E: numpy.ndarray[Any, numpy.dtype[{float64}]]
+
+reveal_type(np.broadcast_to(AR_f8, 5))  # E: numpy.ndarray[Any, numpy.dtype[{float64}]]
+reveal_type(np.broadcast_to(AR_LIKE_f, (1, 5)))  # E: numpy.ndarray[Any, numpy.dtype[Any]]
+reveal_type(np.broadcast_to(AR_f8, [4, 6], subok=True))  # E: numpy.ndarray[Any, numpy.dtype[{float64}]]
+
+reveal_type(np.broadcast_shapes((1, 2), [3, 1], (3, 2)))  # E: tuple[builtins.int]
+reveal_type(np.broadcast_shapes((6, 7), (5, 6, 1), 7, (5, 1, 7)))  # E: tuple[builtins.int]
+
+reveal_type(np.broadcast_arrays(AR_f8, AR_f8))  # E: list[numpy.ndarray[Any, numpy.dtype[Any]]]
+reveal_type(np.broadcast_arrays(AR_f8, AR_LIKE_f))  # E: list[numpy.ndarray[Any, numpy.dtype[Any]]]