twodim_base.pyi
1 from collections.abc import Callable, Sequence 2 from typing import ( 3 Any, 4 overload, 5 TypeVar, 6 Union, 7 ) 8 9 from numpy import ( 10 generic, 11 number, 12 bool_, 13 timedelta64, 14 datetime64, 15 int_, 16 intp, 17 float64, 18 signedinteger, 19 floating, 20 complexfloating, 21 object_, 22 _OrderCF, 23 ) 24 25 from numpy._typing import ( 26 DTypeLike, 27 _DTypeLike, 28 ArrayLike, 29 _ArrayLike, 30 NDArray, 31 _SupportsArrayFunc, 32 _ArrayLikeInt_co, 33 _ArrayLikeFloat_co, 34 _ArrayLikeComplex_co, 35 _ArrayLikeObject_co, 36 ) 37 38 _T = TypeVar("_T") 39 _SCT = TypeVar("_SCT", bound=generic) 40 41 # The returned arrays dtype must be compatible with `np.equal` 42 _MaskFunc = Callable[ 43 [NDArray[int_], _T], 44 NDArray[Union[number[Any], bool_, timedelta64, datetime64, object_]], 45 ] 46 47 __all__: list[str] 48 49 @overload 50 def fliplr(m: _ArrayLike[_SCT]) -> NDArray[_SCT]: ... 51 @overload 52 def fliplr(m: ArrayLike) -> NDArray[Any]: ... 53 54 @overload 55 def flipud(m: _ArrayLike[_SCT]) -> NDArray[_SCT]: ... 56 @overload 57 def flipud(m: ArrayLike) -> NDArray[Any]: ... 58 59 @overload 60 def eye( 61 N: int, 62 M: None | int = ..., 63 k: int = ..., 64 dtype: None = ..., 65 order: _OrderCF = ..., 66 *, 67 like: None | _SupportsArrayFunc = ..., 68 ) -> NDArray[float64]: ... 69 @overload 70 def eye( 71 N: int, 72 M: None | int = ..., 73 k: int = ..., 74 dtype: _DTypeLike[_SCT] = ..., 75 order: _OrderCF = ..., 76 *, 77 like: None | _SupportsArrayFunc = ..., 78 ) -> NDArray[_SCT]: ... 79 @overload 80 def eye( 81 N: int, 82 M: None | int = ..., 83 k: int = ..., 84 dtype: DTypeLike = ..., 85 order: _OrderCF = ..., 86 *, 87 like: None | _SupportsArrayFunc = ..., 88 ) -> NDArray[Any]: ... 89 90 @overload 91 def diag(v: _ArrayLike[_SCT], k: int = ...) -> NDArray[_SCT]: ... 92 @overload 93 def diag(v: ArrayLike, k: int = ...) -> NDArray[Any]: ... 94 95 @overload 96 def diagflat(v: _ArrayLike[_SCT], k: int = ...) -> NDArray[_SCT]: ... 97 @overload 98 def diagflat(v: ArrayLike, k: int = ...) -> NDArray[Any]: ... 99 100 @overload 101 def tri( 102 N: int, 103 M: None | int = ..., 104 k: int = ..., 105 dtype: None = ..., 106 *, 107 like: None | _SupportsArrayFunc = ... 108 ) -> NDArray[float64]: ... 109 @overload 110 def tri( 111 N: int, 112 M: None | int = ..., 113 k: int = ..., 114 dtype: _DTypeLike[_SCT] = ..., 115 *, 116 like: None | _SupportsArrayFunc = ... 117 ) -> NDArray[_SCT]: ... 118 @overload 119 def tri( 120 N: int, 121 M: None | int = ..., 122 k: int = ..., 123 dtype: DTypeLike = ..., 124 *, 125 like: None | _SupportsArrayFunc = ... 126 ) -> NDArray[Any]: ... 127 128 @overload 129 def tril(v: _ArrayLike[_SCT], k: int = ...) -> NDArray[_SCT]: ... 130 @overload 131 def tril(v: ArrayLike, k: int = ...) -> NDArray[Any]: ... 132 133 @overload 134 def triu(v: _ArrayLike[_SCT], k: int = ...) -> NDArray[_SCT]: ... 135 @overload 136 def triu(v: ArrayLike, k: int = ...) -> NDArray[Any]: ... 137 138 @overload 139 def vander( # type: ignore[misc] 140 x: _ArrayLikeInt_co, 141 N: None | int = ..., 142 increasing: bool = ..., 143 ) -> NDArray[signedinteger[Any]]: ... 144 @overload 145 def vander( # type: ignore[misc] 146 x: _ArrayLikeFloat_co, 147 N: None | int = ..., 148 increasing: bool = ..., 149 ) -> NDArray[floating[Any]]: ... 150 @overload 151 def vander( 152 x: _ArrayLikeComplex_co, 153 N: None | int = ..., 154 increasing: bool = ..., 155 ) -> NDArray[complexfloating[Any, Any]]: ... 156 @overload 157 def vander( 158 x: _ArrayLikeObject_co, 159 N: None | int = ..., 160 increasing: bool = ..., 161 ) -> NDArray[object_]: ... 162 163 @overload 164 def histogram2d( # type: ignore[misc] 165 x: _ArrayLikeFloat_co, 166 y: _ArrayLikeFloat_co, 167 bins: int | Sequence[int] = ..., 168 range: None | _ArrayLikeFloat_co = ..., 169 density: None | bool = ..., 170 weights: None | _ArrayLikeFloat_co = ..., 171 ) -> tuple[ 172 NDArray[float64], 173 NDArray[floating[Any]], 174 NDArray[floating[Any]], 175 ]: ... 176 @overload 177 def histogram2d( 178 x: _ArrayLikeComplex_co, 179 y: _ArrayLikeComplex_co, 180 bins: int | Sequence[int] = ..., 181 range: None | _ArrayLikeFloat_co = ..., 182 density: None | bool = ..., 183 weights: None | _ArrayLikeFloat_co = ..., 184 ) -> tuple[ 185 NDArray[float64], 186 NDArray[complexfloating[Any, Any]], 187 NDArray[complexfloating[Any, Any]], 188 ]: ... 189 @overload # TODO: Sort out `bins` 190 def histogram2d( 191 x: _ArrayLikeComplex_co, 192 y: _ArrayLikeComplex_co, 193 bins: Sequence[_ArrayLikeInt_co], 194 range: None | _ArrayLikeFloat_co = ..., 195 density: None | bool = ..., 196 weights: None | _ArrayLikeFloat_co = ..., 197 ) -> tuple[ 198 NDArray[float64], 199 NDArray[Any], 200 NDArray[Any], 201 ]: ... 202 203 # NOTE: we're assuming/demanding here the `mask_func` returns 204 # an ndarray of shape `(n, n)`; otherwise there is the possibility 205 # of the output tuple having more or less than 2 elements 206 @overload 207 def mask_indices( 208 n: int, 209 mask_func: _MaskFunc[int], 210 k: int = ..., 211 ) -> tuple[NDArray[intp], NDArray[intp]]: ... 212 @overload 213 def mask_indices( 214 n: int, 215 mask_func: _MaskFunc[_T], 216 k: _T, 217 ) -> tuple[NDArray[intp], NDArray[intp]]: ... 218 219 def tril_indices( 220 n: int, 221 k: int = ..., 222 m: None | int = ..., 223 ) -> tuple[NDArray[int_], NDArray[int_]]: ... 224 225 def tril_indices_from( 226 arr: NDArray[Any], 227 k: int = ..., 228 ) -> tuple[NDArray[int_], NDArray[int_]]: ... 229 230 def triu_indices( 231 n: int, 232 k: int = ..., 233 m: None | int = ..., 234 ) -> tuple[NDArray[int_], NDArray[int_]]: ... 235 236 def triu_indices_from( 237 arr: NDArray[Any], 238 k: int = ..., 239 ) -> tuple[NDArray[int_], NDArray[int_]]: ...