from typing import Protocol, TypeAlias, TypeVar, assert_type
|
|
import numpy as np
|
from numpy._typing import _64Bit
|
|
_T = TypeVar("_T")
|
_T_co = TypeVar("_T_co", covariant=True)
|
|
class CanAbs(Protocol[_T_co]):
|
def __abs__(self, /) -> _T_co: ...
|
|
class CanInvert(Protocol[_T_co]):
|
def __invert__(self, /) -> _T_co: ...
|
|
class CanNeg(Protocol[_T_co]):
|
def __neg__(self, /) -> _T_co: ...
|
|
class CanPos(Protocol[_T_co]):
|
def __pos__(self, /) -> _T_co: ...
|
|
def do_abs(x: CanAbs[_T]) -> _T: ...
|
def do_invert(x: CanInvert[_T]) -> _T: ...
|
def do_neg(x: CanNeg[_T]) -> _T: ...
|
def do_pos(x: CanPos[_T]) -> _T: ...
|
|
_Bool_1d: TypeAlias = np.ndarray[tuple[int], np.dtype[np.bool]]
|
_UInt8_1d: TypeAlias = np.ndarray[tuple[int], np.dtype[np.uint8]]
|
_Int16_1d: TypeAlias = np.ndarray[tuple[int], np.dtype[np.int16]]
|
_LongLong_1d: TypeAlias = np.ndarray[tuple[int], np.dtype[np.longlong]]
|
_Float32_1d: TypeAlias = np.ndarray[tuple[int], np.dtype[np.float32]]
|
_Float64_1d: TypeAlias = np.ndarray[tuple[int], np.dtype[np.float64]]
|
_LongDouble_1d: TypeAlias = np.ndarray[tuple[int], np.dtype[np.longdouble]]
|
_Complex64_1d: TypeAlias = np.ndarray[tuple[int], np.dtype[np.complex64]]
|
_Complex128_1d: TypeAlias = np.ndarray[tuple[int], np.dtype[np.complex128]]
|
_CLongDouble_1d: TypeAlias = np.ndarray[tuple[int], np.dtype[np.clongdouble]]
|
|
b1_1d: _Bool_1d
|
u1_1d: _UInt8_1d
|
i2_1d: _Int16_1d
|
q_1d: _LongLong_1d
|
f4_1d: _Float32_1d
|
f8_1d: _Float64_1d
|
g_1d: _LongDouble_1d
|
c8_1d: _Complex64_1d
|
c16_1d: _Complex128_1d
|
G_1d: _CLongDouble_1d
|
|
assert_type(do_abs(b1_1d), _Bool_1d)
|
assert_type(do_abs(u1_1d), _UInt8_1d)
|
assert_type(do_abs(i2_1d), _Int16_1d)
|
assert_type(do_abs(q_1d), _LongLong_1d)
|
assert_type(do_abs(f4_1d), _Float32_1d)
|
assert_type(do_abs(f8_1d), _Float64_1d)
|
assert_type(do_abs(g_1d), _LongDouble_1d)
|
|
assert_type(do_abs(c8_1d), _Float32_1d)
|
# NOTE: Unfortunately it's not possible to have this return a `float64` sctype, see
|
# https://github.com/python/mypy/issues/14070
|
assert_type(do_abs(c16_1d), np.ndarray[tuple[int], np.dtype[np.floating[_64Bit]]])
|
assert_type(do_abs(G_1d), _LongDouble_1d)
|
|
assert_type(do_invert(b1_1d), _Bool_1d)
|
assert_type(do_invert(u1_1d), _UInt8_1d)
|
assert_type(do_invert(i2_1d), _Int16_1d)
|
assert_type(do_invert(q_1d), _LongLong_1d)
|
|
assert_type(do_neg(u1_1d), _UInt8_1d)
|
assert_type(do_neg(i2_1d), _Int16_1d)
|
assert_type(do_neg(q_1d), _LongLong_1d)
|
assert_type(do_neg(f4_1d), _Float32_1d)
|
assert_type(do_neg(c16_1d), _Complex128_1d)
|
|
assert_type(do_pos(u1_1d), _UInt8_1d)
|
assert_type(do_pos(i2_1d), _Int16_1d)
|
assert_type(do_pos(q_1d), _LongLong_1d)
|
assert_type(do_pos(f4_1d), _Float32_1d)
|
assert_type(do_pos(c16_1d), _Complex128_1d)
|
