_ufunc.pyi
1 """A module with private type-check-only `numpy.ufunc` subclasses. 2 3 The signatures of the ufuncs are too varied to reasonably type 4 with a single class. So instead, `ufunc` has been expanded into 5 four private subclasses, one for each combination of 6 `~ufunc.nin` and `~ufunc.nout`. 7 8 """ 9 10 from typing import ( 11 Any, 12 Generic, 13 overload, 14 TypeVar, 15 Literal, 16 SupportsIndex, 17 Protocol, 18 ) 19 20 from numpy import ufunc, _CastingKind, _OrderKACF 21 from numpy.typing import NDArray 22 23 from ._shape import _ShapeLike 24 from ._scalars import _ScalarLike_co 25 from ._array_like import ArrayLike, _ArrayLikeBool_co, _ArrayLikeInt_co 26 from ._dtype_like import DTypeLike 27 28 _T = TypeVar("_T") 29 _2Tuple = tuple[_T, _T] 30 _3Tuple = tuple[_T, _T, _T] 31 _4Tuple = tuple[_T, _T, _T, _T] 32 33 _NTypes = TypeVar("_NTypes", bound=int) 34 _IDType = TypeVar("_IDType", bound=Any) 35 _NameType = TypeVar("_NameType", bound=str) 36 37 38 class _SupportsArrayUFunc(Protocol): 39 def __array_ufunc__( 40 self, 41 ufunc: ufunc, 42 method: Literal["__call__", "reduce", "reduceat", "accumulate", "outer", "inner"], 43 *inputs: Any, 44 **kwargs: Any, 45 ) -> Any: ... 46 47 48 # NOTE: In reality `extobj` should be a length of list 3 containing an 49 # int, an int, and a callable, but there's no way to properly express 50 # non-homogenous lists. 51 # Use `Any` over `Union` to avoid issues related to lists invariance. 52 53 # NOTE: `reduce`, `accumulate`, `reduceat` and `outer` raise a ValueError for 54 # ufuncs that don't accept two input arguments and return one output argument. 55 # In such cases the respective methods are simply typed as `None`. 56 57 # NOTE: Similarly, `at` won't be defined for ufuncs that return 58 # multiple outputs; in such cases `at` is typed as `None` 59 60 # NOTE: If 2 output types are returned then `out` must be a 61 # 2-tuple of arrays. Otherwise `None` or a plain array are also acceptable 62 63 class _UFunc_Nin1_Nout1(ufunc, Generic[_NameType, _NTypes, _IDType]): # type: ignore[misc] 64 @property 65 def __name__(self) -> _NameType: ... 66 @property 67 def ntypes(self) -> _NTypes: ... 68 @property 69 def identity(self) -> _IDType: ... 70 @property 71 def nin(self) -> Literal[1]: ... 72 @property 73 def nout(self) -> Literal[1]: ... 74 @property 75 def nargs(self) -> Literal[2]: ... 76 @property 77 def signature(self) -> None: ... 78 @property 79 def reduce(self) -> None: ... 80 @property 81 def accumulate(self) -> None: ... 82 @property 83 def reduceat(self) -> None: ... 84 @property 85 def outer(self) -> None: ... 86 87 @overload 88 def __call__( 89 self, 90 __x1: _ScalarLike_co, 91 out: None = ..., 92 *, 93 where: None | _ArrayLikeBool_co = ..., 94 casting: _CastingKind = ..., 95 order: _OrderKACF = ..., 96 dtype: DTypeLike = ..., 97 subok: bool = ..., 98 signature: str | _2Tuple[None | str] = ..., 99 extobj: list[Any] = ..., 100 ) -> Any: ... 101 @overload 102 def __call__( 103 self, 104 __x1: ArrayLike, 105 out: None | NDArray[Any] | tuple[NDArray[Any]] = ..., 106 *, 107 where: None | _ArrayLikeBool_co = ..., 108 casting: _CastingKind = ..., 109 order: _OrderKACF = ..., 110 dtype: DTypeLike = ..., 111 subok: bool = ..., 112 signature: str | _2Tuple[None | str] = ..., 113 extobj: list[Any] = ..., 114 ) -> NDArray[Any]: ... 115 @overload 116 def __call__( 117 self, 118 __x1: _SupportsArrayUFunc, 119 out: None | NDArray[Any] | tuple[NDArray[Any]] = ..., 120 *, 121 where: None | _ArrayLikeBool_co = ..., 122 casting: _CastingKind = ..., 123 order: _OrderKACF = ..., 124 dtype: DTypeLike = ..., 125 subok: bool = ..., 126 signature: str | _2Tuple[None | str] = ..., 127 extobj: list[Any] = ..., 128 ) -> Any: ... 129 130 def at( 131 self, 132 a: _SupportsArrayUFunc, 133 indices: _ArrayLikeInt_co, 134 /, 135 ) -> None: ... 136 137 class _UFunc_Nin2_Nout1(ufunc, Generic[_NameType, _NTypes, _IDType]): # type: ignore[misc] 138 @property 139 def __name__(self) -> _NameType: ... 140 @property 141 def ntypes(self) -> _NTypes: ... 142 @property 143 def identity(self) -> _IDType: ... 144 @property 145 def nin(self) -> Literal[2]: ... 146 @property 147 def nout(self) -> Literal[1]: ... 148 @property 149 def nargs(self) -> Literal[3]: ... 150 @property 151 def signature(self) -> None: ... 152 153 @overload 154 def __call__( 155 self, 156 __x1: _ScalarLike_co, 157 __x2: _ScalarLike_co, 158 out: None = ..., 159 *, 160 where: None | _ArrayLikeBool_co = ..., 161 casting: _CastingKind = ..., 162 order: _OrderKACF = ..., 163 dtype: DTypeLike = ..., 164 subok: bool = ..., 165 signature: str | _3Tuple[None | str] = ..., 166 extobj: list[Any] = ..., 167 ) -> Any: ... 168 @overload 169 def __call__( 170 self, 171 __x1: ArrayLike, 172 __x2: ArrayLike, 173 out: None | NDArray[Any] | tuple[NDArray[Any]] = ..., 174 *, 175 where: None | _ArrayLikeBool_co = ..., 176 casting: _CastingKind = ..., 177 order: _OrderKACF = ..., 178 dtype: DTypeLike = ..., 179 subok: bool = ..., 180 signature: str | _3Tuple[None | str] = ..., 181 extobj: list[Any] = ..., 182 ) -> NDArray[Any]: ... 183 184 def at( 185 self, 186 a: NDArray[Any], 187 indices: _ArrayLikeInt_co, 188 b: ArrayLike, 189 /, 190 ) -> None: ... 191 192 def reduce( 193 self, 194 array: ArrayLike, 195 axis: None | _ShapeLike = ..., 196 dtype: DTypeLike = ..., 197 out: None | NDArray[Any] = ..., 198 keepdims: bool = ..., 199 initial: Any = ..., 200 where: _ArrayLikeBool_co = ..., 201 ) -> Any: ... 202 203 def accumulate( 204 self, 205 array: ArrayLike, 206 axis: SupportsIndex = ..., 207 dtype: DTypeLike = ..., 208 out: None | NDArray[Any] = ..., 209 ) -> NDArray[Any]: ... 210 211 def reduceat( 212 self, 213 array: ArrayLike, 214 indices: _ArrayLikeInt_co, 215 axis: SupportsIndex = ..., 216 dtype: DTypeLike = ..., 217 out: None | NDArray[Any] = ..., 218 ) -> NDArray[Any]: ... 219 220 # Expand `**kwargs` into explicit keyword-only arguments 221 @overload 222 def outer( 223 self, 224 A: _ScalarLike_co, 225 B: _ScalarLike_co, 226 /, *, 227 out: None = ..., 228 where: None | _ArrayLikeBool_co = ..., 229 casting: _CastingKind = ..., 230 order: _OrderKACF = ..., 231 dtype: DTypeLike = ..., 232 subok: bool = ..., 233 signature: str | _3Tuple[None | str] = ..., 234 extobj: list[Any] = ..., 235 ) -> Any: ... 236 @overload 237 def outer( # type: ignore[misc] 238 self, 239 A: ArrayLike, 240 B: ArrayLike, 241 /, *, 242 out: None | NDArray[Any] | tuple[NDArray[Any]] = ..., 243 where: None | _ArrayLikeBool_co = ..., 244 casting: _CastingKind = ..., 245 order: _OrderKACF = ..., 246 dtype: DTypeLike = ..., 247 subok: bool = ..., 248 signature: str | _3Tuple[None | str] = ..., 249 extobj: list[Any] = ..., 250 ) -> NDArray[Any]: ... 251 252 class _UFunc_Nin1_Nout2(ufunc, Generic[_NameType, _NTypes, _IDType]): # type: ignore[misc] 253 @property 254 def __name__(self) -> _NameType: ... 255 @property 256 def ntypes(self) -> _NTypes: ... 257 @property 258 def identity(self) -> _IDType: ... 259 @property 260 def nin(self) -> Literal[1]: ... 261 @property 262 def nout(self) -> Literal[2]: ... 263 @property 264 def nargs(self) -> Literal[3]: ... 265 @property 266 def signature(self) -> None: ... 267 @property 268 def at(self) -> None: ... 269 @property 270 def reduce(self) -> None: ... 271 @property 272 def accumulate(self) -> None: ... 273 @property 274 def reduceat(self) -> None: ... 275 @property 276 def outer(self) -> None: ... 277 278 @overload 279 def __call__( 280 self, 281 __x1: _ScalarLike_co, 282 __out1: None = ..., 283 __out2: None = ..., 284 *, 285 where: None | _ArrayLikeBool_co = ..., 286 casting: _CastingKind = ..., 287 order: _OrderKACF = ..., 288 dtype: DTypeLike = ..., 289 subok: bool = ..., 290 signature: str | _3Tuple[None | str] = ..., 291 extobj: list[Any] = ..., 292 ) -> _2Tuple[Any]: ... 293 @overload 294 def __call__( 295 self, 296 __x1: ArrayLike, 297 __out1: None | NDArray[Any] = ..., 298 __out2: None | NDArray[Any] = ..., 299 *, 300 out: _2Tuple[NDArray[Any]] = ..., 301 where: None | _ArrayLikeBool_co = ..., 302 casting: _CastingKind = ..., 303 order: _OrderKACF = ..., 304 dtype: DTypeLike = ..., 305 subok: bool = ..., 306 signature: str | _3Tuple[None | str] = ..., 307 extobj: list[Any] = ..., 308 ) -> _2Tuple[NDArray[Any]]: ... 309 @overload 310 def __call__( 311 self, 312 __x1: _SupportsArrayUFunc, 313 __out1: None | NDArray[Any] = ..., 314 __out2: None | NDArray[Any] = ..., 315 *, 316 out: _2Tuple[NDArray[Any]] = ..., 317 where: None | _ArrayLikeBool_co = ..., 318 casting: _CastingKind = ..., 319 order: _OrderKACF = ..., 320 dtype: DTypeLike = ..., 321 subok: bool = ..., 322 signature: str | _3Tuple[None | str] = ..., 323 extobj: list[Any] = ..., 324 ) -> _2Tuple[Any]: ... 325 326 class _UFunc_Nin2_Nout2(ufunc, Generic[_NameType, _NTypes, _IDType]): # type: ignore[misc] 327 @property 328 def __name__(self) -> _NameType: ... 329 @property 330 def ntypes(self) -> _NTypes: ... 331 @property 332 def identity(self) -> _IDType: ... 333 @property 334 def nin(self) -> Literal[2]: ... 335 @property 336 def nout(self) -> Literal[2]: ... 337 @property 338 def nargs(self) -> Literal[4]: ... 339 @property 340 def signature(self) -> None: ... 341 @property 342 def at(self) -> None: ... 343 @property 344 def reduce(self) -> None: ... 345 @property 346 def accumulate(self) -> None: ... 347 @property 348 def reduceat(self) -> None: ... 349 @property 350 def outer(self) -> None: ... 351 352 @overload 353 def __call__( 354 self, 355 __x1: _ScalarLike_co, 356 __x2: _ScalarLike_co, 357 __out1: None = ..., 358 __out2: None = ..., 359 *, 360 where: None | _ArrayLikeBool_co = ..., 361 casting: _CastingKind = ..., 362 order: _OrderKACF = ..., 363 dtype: DTypeLike = ..., 364 subok: bool = ..., 365 signature: str | _4Tuple[None | str] = ..., 366 extobj: list[Any] = ..., 367 ) -> _2Tuple[Any]: ... 368 @overload 369 def __call__( 370 self, 371 __x1: ArrayLike, 372 __x2: ArrayLike, 373 __out1: None | NDArray[Any] = ..., 374 __out2: None | NDArray[Any] = ..., 375 *, 376 out: _2Tuple[NDArray[Any]] = ..., 377 where: None | _ArrayLikeBool_co = ..., 378 casting: _CastingKind = ..., 379 order: _OrderKACF = ..., 380 dtype: DTypeLike = ..., 381 subok: bool = ..., 382 signature: str | _4Tuple[None | str] = ..., 383 extobj: list[Any] = ..., 384 ) -> _2Tuple[NDArray[Any]]: ... 385 386 class _GUFunc_Nin2_Nout1(ufunc, Generic[_NameType, _NTypes, _IDType]): # type: ignore[misc] 387 @property 388 def __name__(self) -> _NameType: ... 389 @property 390 def ntypes(self) -> _NTypes: ... 391 @property 392 def identity(self) -> _IDType: ... 393 @property 394 def nin(self) -> Literal[2]: ... 395 @property 396 def nout(self) -> Literal[1]: ... 397 @property 398 def nargs(self) -> Literal[3]: ... 399 400 # NOTE: In practice the only gufunc in the main namespace is `matmul`, 401 # so we can use its signature here 402 @property 403 def signature(self) -> Literal["(n?,k),(k,m?)->(n?,m?)"]: ... 404 @property 405 def reduce(self) -> None: ... 406 @property 407 def accumulate(self) -> None: ... 408 @property 409 def reduceat(self) -> None: ... 410 @property 411 def outer(self) -> None: ... 412 @property 413 def at(self) -> None: ... 414 415 # Scalar for 1D array-likes; ndarray otherwise 416 @overload 417 def __call__( 418 self, 419 __x1: ArrayLike, 420 __x2: ArrayLike, 421 out: None = ..., 422 *, 423 casting: _CastingKind = ..., 424 order: _OrderKACF = ..., 425 dtype: DTypeLike = ..., 426 subok: bool = ..., 427 signature: str | _3Tuple[None | str] = ..., 428 extobj: list[Any] = ..., 429 axes: list[_2Tuple[SupportsIndex]] = ..., 430 ) -> Any: ... 431 @overload 432 def __call__( 433 self, 434 __x1: ArrayLike, 435 __x2: ArrayLike, 436 out: NDArray[Any] | tuple[NDArray[Any]], 437 *, 438 casting: _CastingKind = ..., 439 order: _OrderKACF = ..., 440 dtype: DTypeLike = ..., 441 subok: bool = ..., 442 signature: str | _3Tuple[None | str] = ..., 443 extobj: list[Any] = ..., 444 axes: list[_2Tuple[SupportsIndex]] = ..., 445 ) -> NDArray[Any]: ...