shape_base.pyi
1 from collections.abc import Callable, Sequence 2 from typing import TypeVar, Any, overload, SupportsIndex, Protocol 3 4 from numpy import ( 5 generic, 6 integer, 7 ufunc, 8 bool_, 9 unsignedinteger, 10 signedinteger, 11 floating, 12 complexfloating, 13 object_, 14 ) 15 16 from numpy._typing import ( 17 ArrayLike, 18 NDArray, 19 _ShapeLike, 20 _ArrayLike, 21 _ArrayLikeBool_co, 22 _ArrayLikeUInt_co, 23 _ArrayLikeInt_co, 24 _ArrayLikeFloat_co, 25 _ArrayLikeComplex_co, 26 _ArrayLikeObject_co, 27 ) 28 29 from numpy.core.shape_base import vstack 30 31 _SCT = TypeVar("_SCT", bound=generic) 32 33 # The signatures of `__array_wrap__` and `__array_prepare__` are the same; 34 # give them unique names for the sake of clarity 35 class _ArrayWrap(Protocol): 36 def __call__( 37 self, 38 array: NDArray[Any], 39 context: None | tuple[ufunc, tuple[Any, ...], int] = ..., 40 /, 41 ) -> Any: ... 42 43 class _ArrayPrepare(Protocol): 44 def __call__( 45 self, 46 array: NDArray[Any], 47 context: None | tuple[ufunc, tuple[Any, ...], int] = ..., 48 /, 49 ) -> Any: ... 50 51 class _SupportsArrayWrap(Protocol): 52 @property 53 def __array_wrap__(self) -> _ArrayWrap: ... 54 55 class _SupportsArrayPrepare(Protocol): 56 @property 57 def __array_prepare__(self) -> _ArrayPrepare: ... 58 59 __all__: list[str] 60 61 row_stack = vstack 62 63 def take_along_axis( 64 arr: _SCT | NDArray[_SCT], 65 indices: NDArray[integer[Any]], 66 axis: None | int, 67 ) -> NDArray[_SCT]: ... 68 69 def put_along_axis( 70 arr: NDArray[_SCT], 71 indices: NDArray[integer[Any]], 72 values: ArrayLike, 73 axis: None | int, 74 ) -> None: ... 75 76 # TODO: Use PEP 612 `ParamSpec` once mypy supports `Concatenate` 77 # xref python/mypy#8645 78 @overload 79 def apply_along_axis( 80 func1d: Callable[..., _ArrayLike[_SCT]], 81 axis: SupportsIndex, 82 arr: ArrayLike, 83 *args: Any, 84 **kwargs: Any, 85 ) -> NDArray[_SCT]: ... 86 @overload 87 def apply_along_axis( 88 func1d: Callable[..., ArrayLike], 89 axis: SupportsIndex, 90 arr: ArrayLike, 91 *args: Any, 92 **kwargs: Any, 93 ) -> NDArray[Any]: ... 94 95 def apply_over_axes( 96 func: Callable[[NDArray[Any], int], NDArray[_SCT]], 97 a: ArrayLike, 98 axes: int | Sequence[int], 99 ) -> NDArray[_SCT]: ... 100 101 @overload 102 def expand_dims( 103 a: _ArrayLike[_SCT], 104 axis: _ShapeLike, 105 ) -> NDArray[_SCT]: ... 106 @overload 107 def expand_dims( 108 a: ArrayLike, 109 axis: _ShapeLike, 110 ) -> NDArray[Any]: ... 111 112 @overload 113 def column_stack(tup: Sequence[_ArrayLike[_SCT]]) -> NDArray[_SCT]: ... 114 @overload 115 def column_stack(tup: Sequence[ArrayLike]) -> NDArray[Any]: ... 116 117 @overload 118 def dstack(tup: Sequence[_ArrayLike[_SCT]]) -> NDArray[_SCT]: ... 119 @overload 120 def dstack(tup: Sequence[ArrayLike]) -> NDArray[Any]: ... 121 122 @overload 123 def array_split( 124 ary: _ArrayLike[_SCT], 125 indices_or_sections: _ShapeLike, 126 axis: SupportsIndex = ..., 127 ) -> list[NDArray[_SCT]]: ... 128 @overload 129 def array_split( 130 ary: ArrayLike, 131 indices_or_sections: _ShapeLike, 132 axis: SupportsIndex = ..., 133 ) -> list[NDArray[Any]]: ... 134 135 @overload 136 def split( 137 ary: _ArrayLike[_SCT], 138 indices_or_sections: _ShapeLike, 139 axis: SupportsIndex = ..., 140 ) -> list[NDArray[_SCT]]: ... 141 @overload 142 def split( 143 ary: ArrayLike, 144 indices_or_sections: _ShapeLike, 145 axis: SupportsIndex = ..., 146 ) -> list[NDArray[Any]]: ... 147 148 @overload 149 def hsplit( 150 ary: _ArrayLike[_SCT], 151 indices_or_sections: _ShapeLike, 152 ) -> list[NDArray[_SCT]]: ... 153 @overload 154 def hsplit( 155 ary: ArrayLike, 156 indices_or_sections: _ShapeLike, 157 ) -> list[NDArray[Any]]: ... 158 159 @overload 160 def vsplit( 161 ary: _ArrayLike[_SCT], 162 indices_or_sections: _ShapeLike, 163 ) -> list[NDArray[_SCT]]: ... 164 @overload 165 def vsplit( 166 ary: ArrayLike, 167 indices_or_sections: _ShapeLike, 168 ) -> list[NDArray[Any]]: ... 169 170 @overload 171 def dsplit( 172 ary: _ArrayLike[_SCT], 173 indices_or_sections: _ShapeLike, 174 ) -> list[NDArray[_SCT]]: ... 175 @overload 176 def dsplit( 177 ary: ArrayLike, 178 indices_or_sections: _ShapeLike, 179 ) -> list[NDArray[Any]]: ... 180 181 @overload 182 def get_array_prepare(*args: _SupportsArrayPrepare) -> _ArrayPrepare: ... 183 @overload 184 def get_array_prepare(*args: object) -> None | _ArrayPrepare: ... 185 186 @overload 187 def get_array_wrap(*args: _SupportsArrayWrap) -> _ArrayWrap: ... 188 @overload 189 def get_array_wrap(*args: object) -> None | _ArrayWrap: ... 190 191 @overload 192 def kron(a: _ArrayLikeBool_co, b: _ArrayLikeBool_co) -> NDArray[bool_]: ... # type: ignore[misc] 193 @overload 194 def kron(a: _ArrayLikeUInt_co, b: _ArrayLikeUInt_co) -> NDArray[unsignedinteger[Any]]: ... # type: ignore[misc] 195 @overload 196 def kron(a: _ArrayLikeInt_co, b: _ArrayLikeInt_co) -> NDArray[signedinteger[Any]]: ... # type: ignore[misc] 197 @overload 198 def kron(a: _ArrayLikeFloat_co, b: _ArrayLikeFloat_co) -> NDArray[floating[Any]]: ... # type: ignore[misc] 199 @overload 200 def kron(a: _ArrayLikeComplex_co, b: _ArrayLikeComplex_co) -> NDArray[complexfloating[Any, Any]]: ... 201 @overload 202 def kron(a: _ArrayLikeObject_co, b: Any) -> NDArray[object_]: ... 203 @overload 204 def kron(a: Any, b: _ArrayLikeObject_co) -> NDArray[object_]: ... 205 206 @overload 207 def tile( 208 A: _ArrayLike[_SCT], 209 reps: int | Sequence[int], 210 ) -> NDArray[_SCT]: ... 211 @overload 212 def tile( 213 A: ArrayLike, 214 reps: int | Sequence[int], 215 ) -> NDArray[Any]: ...