diff options
Diffstat (limited to 'numpy')
| -rw-r--r-- | numpy/lib/stride_tricks.pyi | 84 | ||||
| -rw-r--r-- | numpy/typing/tests/data/fail/stride_tricks.py | 9 | ||||
| -rw-r--r-- | numpy/typing/tests/data/reveal/stride_tricks.py | 28 |
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]]] |
