summaryrefslogtreecommitdiff
path: root/numpy
diff options
context:
space:
mode:
Diffstat (limited to 'numpy')
-rw-r--r--numpy/lib/stride_tricks.pyi84
-rw-r--r--numpy/typing/tests/data/fail/stride_tricks.py9
-rw-r--r--numpy/typing/tests/data/reveal/stride_tricks.py28
3 files changed, 112 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]]: ...
diff --git a/numpy/typing/tests/data/fail/stride_tricks.py b/numpy/typing/tests/data/fail/stride_tricks.py
new file mode 100644
index 000000000..f2bfba743
--- /dev/null
+++ b/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
diff --git a/numpy/typing/tests/data/reveal/stride_tricks.py b/numpy/typing/tests/data/reveal/stride_tricks.py
new file mode 100644
index 000000000..152d9cea6
--- /dev/null
+++ b/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]]]