446 lines
12 KiB
Python
446 lines
12 KiB
Python
"""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]: ...
|