diff options
Diffstat (limited to 'numpy/typing/_callable.py')
| -rw-r--r-- | numpy/typing/_callable.py | 124 |
1 files changed, 123 insertions, 1 deletions
diff --git a/numpy/typing/_callable.py b/numpy/typing/_callable.py index b16891b0e..8b7bfedac 100644 --- a/numpy/typing/_callable.py +++ b/numpy/typing/_callable.py @@ -9,7 +9,15 @@ See the `Mypy documentation`_ on protocols for more details. """ import sys -from typing import Union, TypeVar, overload, Any, TYPE_CHECKING, NoReturn +from typing import ( + Union, + TypeVar, + overload, + Any, + Tuple, + NoReturn, + TYPE_CHECKING, +) from numpy import ( generic, @@ -19,6 +27,7 @@ from numpy import ( integer, unsignedinteger, signedinteger, + int8, int32, int64, floating, @@ -49,9 +58,14 @@ else: HAVE_PROTOCOL = True if TYPE_CHECKING or HAVE_PROTOCOL: + _T = TypeVar("_T") + _2Tuple = Tuple[_T, _T] + _NBit_co = TypeVar("_NBit_co", covariant=True, bound=NBitBase) _NBit = TypeVar("_NBit", bound=NBitBase) + _IntType = TypeVar("_IntType", bound=integer) + _FloatType = TypeVar("_FloatType", bound=floating) _NumberType = TypeVar("_NumberType", bound=number) _NumberType_co = TypeVar("_NumberType_co", covariant=True, bound=number) _GenericType_co = TypeVar("_GenericType_co", covariant=True, bound=generic) @@ -97,6 +111,30 @@ if TYPE_CHECKING or HAVE_PROTOCOL: @overload def __call__(self, __other: _NumberType) -> _NumberType: ... + class _BoolMod(Protocol): + @overload + def __call__(self, __other: _BoolLike) -> int8: ... + @overload # platform dependent + def __call__(self, __other: int) -> signedinteger[Any]: ... + @overload + def __call__(self, __other: float) -> float64: ... + @overload + def __call__(self, __other: _IntType) -> _IntType: ... + @overload + def __call__(self, __other: _FloatType) -> _FloatType: ... + + class _BoolDivMod(Protocol): + @overload + def __call__(self, __other: _BoolLike) -> _2Tuple[int8]: ... + @overload # platform dependent + def __call__(self, __other: int) -> _2Tuple[signedinteger[Any]]: ... + @overload + def __call__(self, __other: float) -> _2Tuple[float64]: ... + @overload + def __call__(self, __other: _IntType) -> _2Tuple[_IntType]: ... + @overload + def __call__(self, __other: _FloatType) -> _2Tuple[_FloatType]: ... + class _TD64Div(Protocol[_NumberType_co]): @overload def __call__(self, __other: timedelta64) -> _NumberType_co: ... @@ -144,6 +182,34 @@ if TYPE_CHECKING or HAVE_PROTOCOL: self, __other: unsignedinteger[_NBit] ) -> unsignedinteger[Union[_NBit_co, _NBit]]: ... + class _UnsignedIntMod(Protocol[_NBit_co]): + @overload + def __call__(self, __other: bool) -> unsignedinteger[_NBit_co]: ... + @overload + def __call__( + self, __other: Union[int, signedinteger[Any]] + ) -> Union[signedinteger[Any], float64]: ... + @overload + def __call__(self, __other: float) -> float64: ... + @overload + def __call__( + self, __other: unsignedinteger[_NBit] + ) -> unsignedinteger[Union[_NBit_co, _NBit]]: ... + + class _UnsignedIntDivMod(Protocol[_NBit_co]): + @overload + def __call__(self, __other: bool) -> _2Tuple[signedinteger[_NBit_co]]: ... + @overload + def __call__( + self, __other: Union[int, signedinteger[Any]] + ) -> Union[_2Tuple[signedinteger[Any]], _2Tuple[float64]]: ... + @overload + def __call__(self, __other: float) -> _2Tuple[float64]: ... + @overload + def __call__( + self, __other: unsignedinteger[_NBit] + ) -> _2Tuple[unsignedinteger[Union[_NBit_co, _NBit]]]: ... + class _SignedIntOp(Protocol[_NBit_co]): @overload def __call__(self, __other: bool) -> signedinteger[_NBit_co]: ... @@ -168,6 +234,30 @@ if TYPE_CHECKING or HAVE_PROTOCOL: self, __other: signedinteger[_NBit] ) -> signedinteger[Union[_NBit_co, _NBit]]: ... + class _SignedIntMod(Protocol[_NBit_co]): + @overload + def __call__(self, __other: bool) -> signedinteger[_NBit_co]: ... + @overload + def __call__(self, __other: int) -> signedinteger[Any]: ... + @overload + def __call__(self, __other: float) -> float64: ... + @overload + def __call__( + self, __other: signedinteger[_NBit] + ) -> signedinteger[Union[_NBit_co, _NBit]]: ... + + class _SignedIntDivMod(Protocol[_NBit_co]): + @overload + def __call__(self, __other: bool) -> _2Tuple[signedinteger[_NBit_co]]: ... + @overload + def __call__(self, __other: int) -> _2Tuple[signedinteger[Any]]: ... + @overload + def __call__(self, __other: float) -> _2Tuple[float64]: ... + @overload + def __call__( + self, __other: signedinteger[_NBit] + ) -> _2Tuple[signedinteger[Union[_NBit_co, _NBit]]]: ... + class _FloatOp(Protocol[_NBit_co]): @overload def __call__(self, __other: bool) -> floating[_NBit_co]: ... @@ -182,6 +272,30 @@ if TYPE_CHECKING or HAVE_PROTOCOL: self, __other: Union[integer[_NBit], floating[_NBit]] ) -> floating[Union[_NBit_co, _NBit]]: ... + class _FloatMod(Protocol[_NBit_co]): + @overload + def __call__(self, __other: bool) -> floating[_NBit_co]: ... + @overload + def __call__(self, __other: int) -> floating[Any]: ... + @overload + def __call__(self, __other: float) -> float64: ... + @overload + def __call__( + self, __other: Union[integer[_NBit], floating[_NBit]] + ) -> floating[Union[_NBit_co, _NBit]]: ... + + class _FloatDivMod(Protocol[_NBit_co]): + @overload + def __call__(self, __other: bool) -> _2Tuple[floating[_NBit_co]]: ... + @overload + def __call__(self, __other: int) -> _2Tuple[floating[Any]]: ... + @overload + def __call__(self, __other: float) -> _2Tuple[float64]: ... + @overload + def __call__( + self, __other: Union[integer[_NBit], floating[_NBit]] + ) -> _2Tuple[floating[Union[_NBit_co, _NBit]]]: ... + class _ComplexOp(Protocol[_NBit_co]): @overload def __call__(self, __other: bool) -> complexfloating[_NBit_co, _NBit_co]: ... @@ -207,12 +321,20 @@ else: _BoolBitOp = Any _BoolSub = Any _BoolTrueDiv = Any + _BoolMod = Any + _BoolDivMod = Any _TD64Div = Any _IntTrueDiv = Any _UnsignedIntOp = Any _UnsignedIntBitOp = Any + _UnsignedIntMod = Any + _UnsignedIntDivMod = Any _SignedIntOp = Any _SignedIntBitOp = Any + _SignedIntMod = Any + _SignedIntDivMod = Any _FloatOp = Any + _FloatMod = Any + _FloatDivMod = Any _ComplexOp = Any _NumberOp = Any |