JezK
Edit File: _ufunc.pyi
"""A module with private type-check-only `numpy.ufunc` subclasses. The signatures of the ufuncs are too varied to reasonably type with a single class. So instead, `ufunc` has been expanded into four private subclasses, one for each combination of `~ufunc.nin` and `~ufunc.nout`. """ from typing import ( Any, Generic, overload, TypeVar, Literal, SupportsIndex, Protocol, ) from numpy import ufunc, _CastingKind, _OrderKACF from numpy.typing import NDArray from ._shape import _ShapeLike from ._scalars import _ScalarLike_co from ._array_like import ArrayLike, _ArrayLikeBool_co, _ArrayLikeInt_co from ._dtype_like import DTypeLike _T = TypeVar("_T") _2Tuple = tuple[_T, _T] _3Tuple = tuple[_T, _T, _T] _4Tuple = tuple[_T, _T, _T, _T] _NTypes = TypeVar("_NTypes", bound=int) _IDType = TypeVar("_IDType", bound=Any) _NameType = TypeVar("_NameType", bound=str) class _SupportsArrayUFunc(Protocol): def __array_ufunc__( self, ufunc: ufunc, method: Literal["__call__", "reduce", "reduceat", "accumulate", "outer", "inner"], *inputs: Any, **kwargs: Any, ) -> Any: ... # NOTE: In reality `extobj` should be a length of list 3 containing an # int, an int, and a callable, but there's no way to properly express # non-homogenous lists. # Use `Any` over `Union` to avoid issues related to lists invariance. # NOTE: `reduce`, `accumulate`, `reduceat` and `outer` raise a ValueError for # ufuncs that don't accept two input arguments and return one output argument. # In such cases the respective methods are simply typed as `None`. # NOTE: Similarly, `at` won't be defined for ufuncs that return # multiple outputs; in such cases `at` is typed as `None` # NOTE: If 2 output types are returned then `out` must be a # 2-tuple of arrays. Otherwise `None` or a plain array are also acceptable class _UFunc_Nin1_Nout1(ufunc, Generic[_NameType, _NTypes, _IDType]): # type: ignore[misc] @property def __name__(self) -> _NameType: ... @property def ntypes(self) -> _NTypes: ... @property def identity(self) -> _IDType: ... @property def nin(self) -> Literal[1]: ... @property def nout(self) -> Literal[1]: ... @property def nargs(self) -> Literal[2]: ... @property def signature(self) -> None: ... @property def reduce(self) -> None: ... @property def accumulate(self) -> None: ... @property def reduceat(self) -> None: ... @property def outer(self) -> None: ... @overload def __call__( self, __x1: _ScalarLike_co, out: None = ..., *, where: None | _ArrayLikeBool_co = ..., casting: _CastingKind = ..., order: _OrderKACF = ..., dtype: DTypeLike = ..., subok: bool = ..., signature: str | _2Tuple[None | str] = ..., extobj: list[Any] = ..., ) -> Any: ... @overload def __call__( self, __x1: ArrayLike, out: None | NDArray[Any] | tuple[NDArray[Any]] = ..., *, where: None | _ArrayLikeBool_co = ..., casting: _CastingKind = ..., order: _OrderKACF = ..., dtype: DTypeLike = ..., subok: bool = ..., signature: str | _2Tuple[None | str] = ..., extobj: list[Any] = ..., ) -> NDArray[Any]: ... @overload def __call__( self, __x1: _SupportsArrayUFunc, out: None | NDArray[Any] | tuple[NDArray[Any]] = ..., *, where: None | _ArrayLikeBool_co = ..., casting: _CastingKind = ..., order: _OrderKACF = ..., dtype: DTypeLike = ..., subok: bool = ..., signature: str | _2Tuple[None | str] = ..., extobj: list[Any] = ..., ) -> Any: ... def at( self, a: _SupportsArrayUFunc, indices: _ArrayLikeInt_co, /, ) -> None: ... class _UFunc_Nin2_Nout1(ufunc, Generic[_NameType, _NTypes, _IDType]): # type: ignore[misc] @property def __name__(self) -> _NameType: ... @property def ntypes(self) -> _NTypes: ... @property def identity(self) -> _IDType: ... @property def nin(self) -> Literal[2]: ... @property def nout(self) -> Literal[1]: ... @property def nargs(self) -> Literal[3]: ... @property def signature(self) -> None: ... @overload def __call__( self, __x1: _ScalarLike_co, __x2: _ScalarLike_co, out: None = ..., *, where: None | _ArrayLikeBool_co = ..., casting: _CastingKind = ..., order: _OrderKACF = ..., dtype: DTypeLike = ..., subok: bool = ..., signature: str | _3Tuple[None | str] = ..., extobj: list[Any] = ..., ) -> Any: ... @overload def __call__( self, __x1: ArrayLike, __x2: ArrayLike, out: None | NDArray[Any] | tuple[NDArray[Any]] = ..., *, where: None | _ArrayLikeBool_co = ..., casting: _CastingKind = ..., order: _OrderKACF = ..., dtype: DTypeLike = ..., subok: bool = ..., signature: str | _3Tuple[None | str] = ..., extobj: list[Any] = ..., ) -> NDArray[Any]: ... def at( self, a: NDArray[Any], indices: _ArrayLikeInt_co, b: ArrayLike, /, ) -> None: ... def reduce( self, array: ArrayLike, axis: None | _ShapeLike = ..., dtype: DTypeLike = ..., out: None | NDArray[Any] = ..., keepdims: bool = ..., initial: Any = ..., where: _ArrayLikeBool_co = ..., ) -> Any: ... def accumulate( self, array: ArrayLike, axis: SupportsIndex = ..., dtype: DTypeLike = ..., out: None | NDArray[Any] = ..., ) -> NDArray[Any]: ... def reduceat( self, array: ArrayLike, indices: _ArrayLikeInt_co, axis: SupportsIndex = ..., dtype: DTypeLike = ..., out: None | NDArray[Any] = ..., ) -> NDArray[Any]: ... # Expand `**kwargs` into explicit keyword-only arguments @overload def outer( self, A: _ScalarLike_co, B: _ScalarLike_co, /, *, out: None = ..., where: None | _ArrayLikeBool_co = ..., casting: _CastingKind = ..., order: _OrderKACF = ..., dtype: DTypeLike = ..., subok: bool = ..., signature: str | _3Tuple[None | str] = ..., extobj: list[Any] = ..., ) -> Any: ... @overload def outer( # type: ignore[misc] self, A: ArrayLike, B: ArrayLike, /, *, out: None | NDArray[Any] | tuple[NDArray[Any]] = ..., where: None | _ArrayLikeBool_co = ..., casting: _CastingKind = ..., order: _OrderKACF = ..., dtype: DTypeLike = ..., subok: bool = ..., signature: str | _3Tuple[None | str] = ..., extobj: list[Any] = ..., ) -> NDArray[Any]: ... class _UFunc_Nin1_Nout2(ufunc, Generic[_NameType, _NTypes, _IDType]): # type: ignore[misc] @property def __name__(self) -> _NameType: ... @property def ntypes(self) -> _NTypes: ... @property def identity(self) -> _IDType: ... @property def nin(self) -> Literal[1]: ... @property def nout(self) -> Literal[2]: ... @property def nargs(self) -> Literal[3]: ... @property def signature(self) -> None: ... @property def at(self) -> None: ... @property def reduce(self) -> None: ... @property def accumulate(self) -> None: ... @property def reduceat(self) -> None: ... @property def outer(self) -> None: ... @overload def __call__( self, __x1: _ScalarLike_co, __out1: None = ..., __out2: None = ..., *, where: None | _ArrayLikeBool_co = ..., casting: _CastingKind = ..., order: _OrderKACF = ..., dtype: DTypeLike = ..., subok: bool = ..., signature: str | _3Tuple[None | str] = ..., extobj: list[Any] = ..., ) -> _2Tuple[Any]: ... @overload def __call__( self, __x1: ArrayLike, __out1: None | NDArray[Any] = ..., __out2: None | NDArray[Any] = ..., *, out: _2Tuple[NDArray[Any]] = ..., where: None | _ArrayLikeBool_co = ..., casting: _CastingKind = ..., order: _OrderKACF = ..., dtype: DTypeLike = ..., subok: bool = ..., signature: str | _3Tuple[None | str] = ..., extobj: list[Any] = ..., ) -> _2Tuple[NDArray[Any]]: ... @overload def __call__( self, __x1: _SupportsArrayUFunc, __out1: None | NDArray[Any] = ..., __out2: None | NDArray[Any] = ..., *, out: _2Tuple[NDArray[Any]] = ..., where: None | _ArrayLikeBool_co = ..., casting: _CastingKind = ..., order: _OrderKACF = ..., dtype: DTypeLike = ..., subok: bool = ..., signature: str | _3Tuple[None | str] = ..., extobj: list[Any] = ..., ) -> _2Tuple[Any]: ... class _UFunc_Nin2_Nout2(ufunc, Generic[_NameType, _NTypes, _IDType]): # type: ignore[misc] @property def __name__(self) -> _NameType: ... @property def ntypes(self) -> _NTypes: ... @property def identity(self) -> _IDType: ... @property def nin(self) -> Literal[2]: ... @property def nout(self) -> Literal[2]: ... @property def nargs(self) -> Literal[4]: ... @property def signature(self) -> None: ... @property def at(self) -> None: ... @property def reduce(self) -> None: ... @property def accumulate(self) -> None: ... @property def reduceat(self) -> None: ... @property def outer(self) -> None: ... @overload def __call__( self, __x1: _ScalarLike_co, __x2: _ScalarLike_co, __out1: None = ..., __out2: None = ..., *, where: None | _ArrayLikeBool_co = ..., casting: _CastingKind = ..., order: _OrderKACF = ..., dtype: DTypeLike = ..., subok: bool = ..., signature: str | _4Tuple[None | str] = ..., extobj: list[Any] = ..., ) -> _2Tuple[Any]: ... @overload def __call__( self, __x1: ArrayLike, __x2: ArrayLike, __out1: None | NDArray[Any] = ..., __out2: None | NDArray[Any] = ..., *, out: _2Tuple[NDArray[Any]] = ..., where: None | _ArrayLikeBool_co = ..., casting: _CastingKind = ..., order: _OrderKACF = ..., dtype: DTypeLike = ..., subok: bool = ..., signature: str | _4Tuple[None | str] = ..., extobj: list[Any] = ..., ) -> _2Tuple[NDArray[Any]]: ... class _GUFunc_Nin2_Nout1(ufunc, Generic[_NameType, _NTypes, _IDType]): # type: ignore[misc] @property def __name__(self) -> _NameType: ... @property def ntypes(self) -> _NTypes: ... @property def identity(self) -> _IDType: ... @property def nin(self) -> Literal[2]: ... @property def nout(self) -> Literal[1]: ... @property def nargs(self) -> Literal[3]: ... # NOTE: In practice the only gufunc in the main namespace is `matmul`, # so we can use its signature here @property def signature(self) -> Literal["(n?,k),(k,m?)->(n?,m?)"]: ... @property def reduce(self) -> None: ... @property def accumulate(self) -> None: ... @property def reduceat(self) -> None: ... @property def outer(self) -> None: ... @property def at(self) -> None: ... # Scalar for 1D array-likes; ndarray otherwise @overload def __call__( self, __x1: ArrayLike, __x2: ArrayLike, out: None = ..., *, casting: _CastingKind = ..., order: _OrderKACF = ..., dtype: DTypeLike = ..., subok: bool = ..., signature: str | _3Tuple[None | str] = ..., extobj: list[Any] = ..., axes: list[_2Tuple[SupportsIndex]] = ..., ) -> Any: ... @overload def __call__( self, __x1: ArrayLike, __x2: ArrayLike, out: NDArray[Any] | tuple[NDArray[Any]], *, casting: _CastingKind = ..., order: _OrderKACF = ..., dtype: DTypeLike = ..., subok: bool = ..., signature: str | _3Tuple[None | str] = ..., extobj: list[Any] = ..., axes: list[_2Tuple[SupportsIndex]] = ..., ) -> NDArray[Any]: ...