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

1 files changed, 75 insertions, 9 deletions

diff --git a/numpy/lib/stride_tricks.pyi b/numpy/lib/stride_tricks.pyi
index d2e744b5a..9e4e46b8b 100644
--- a/numpy/lib/stride_tricks.pyi
+++ b/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]]: ...