Merge pull request numpy#27659 from jorenham/typing/transparent-ndarray-ops

TYP: Transparent ``ndarray`` unary operator method signatures

Author: charris

numpy/__init__.pyi

@@ -1760,11 +1760,15 @@ _DType = TypeVar("_DType", bound=dtype[Any])
 _DType_co = TypeVar("_DType_co", covariant=True, bound=dtype[Any])
 _FlexDType = TypeVar("_FlexDType", bound=dtype[flexible])
 
+_IntegralArrayT = TypeVar("_IntegralArrayT", bound=NDArray[integer[Any] | np.bool | object_])
+_RealArrayT = TypeVar("_RealArrayT", bound=NDArray[floating[Any] | integer[Any] | timedelta64 | np.bool | object_])
+_NumericArrayT = TypeVar("_NumericArrayT", bound=NDArray[number[Any] | timedelta64 | object_])
+
 _Shape1D: TypeAlias = tuple[int]
 _Shape2D: TypeAlias = tuple[int, int]
 
+_ShapeType = TypeVar("_ShapeType", bound=_Shape)
 _ShapeType_co = TypeVar("_ShapeType_co", covariant=True, bound=_Shape)
-_ShapeType2 = TypeVar("_ShapeType2", bound=_Shape)
 _Shape2DType_co = TypeVar("_Shape2DType_co", covariant=True, bound=_Shape2D)
 _NumberType = TypeVar("_NumberType", bound=number[Any])
 
@@ -1881,11 +1885,11 @@ class ndarray(_ArrayOrScalarCommon, Generic[_ShapeType_co, _DType_co]):
 
     def __array_wrap__(
         self,
-        array: ndarray[_ShapeType2, _DType],
+        array: ndarray[_ShapeType, _DType],
         context: None | tuple[ufunc, tuple[Any, ...], int] = ...,
         return_scalar: builtins.bool = ...,
         /,
-    ) -> ndarray[_ShapeType2, _DType]: ...
+    ) -> ndarray[_ShapeType, _DType]: ...
 
     @overload
     def __getitem__(self, key: (
@@ -2237,22 +2241,10 @@ class ndarray(_ArrayOrScalarCommon, Generic[_ShapeType_co, _DType_co]):
         offset: SupportsIndex = ...
     ) -> NDArray[Any]: ...
 
-    # Dispatch to the underlying `generic` via protocols
-    def __int__(
-        self: NDArray[SupportsInt],  # type: ignore[type-var]
-    ) -> int: ...
-
-    def __float__(
-        self: NDArray[SupportsFloat],  # type: ignore[type-var]
-    ) -> float: ...
-
-    def __complex__(
-        self: NDArray[SupportsComplex],  # type: ignore[type-var]
-    ) -> complex: ...
-
-    def __index__(
-        self: NDArray[SupportsIndex],  # type: ignore[type-var]
-    ) -> int: ...
+    def __index__(self: NDArray[np.integer[Any]], /) -> int: ...
+    def __int__(self: NDArray[number[Any] | np.bool | object_], /) -> int: ...
+    def __float__(self: NDArray[number[Any] | np.bool | object_], /) -> float: ...
+    def __complex__(self: NDArray[number[Any] | np.bool | object_], /) -> complex: ...
 
     def __len__(self) -> int: ...
     def __setitem__(self, key, value): ...
@@ -2310,41 +2302,25 @@ class ndarray(_ArrayOrScalarCommon, Generic[_ShapeType_co, _DType_co]):
     def __ge__(self: NDArray[Any], other: _ArrayLikeObject_co, /) -> NDArray[np.bool]: ...
 
     # Unary ops
-    @overload
-    def __abs__(self: NDArray[_UnknownType]) -> NDArray[Any]: ...
-    @overload
-    def __abs__(self: NDArray[np.bool]) -> NDArray[np.bool]: ...
-    @overload
-    def __abs__(self: NDArray[complexfloating[_NBit1, _NBit1]]) -> NDArray[floating[_NBit1]]: ...
-    @overload
-    def __abs__(self: NDArray[_NumberType]) -> NDArray[_NumberType]: ...
-    @overload
-    def __abs__(self: NDArray[timedelta64]) -> NDArray[timedelta64]: ...
-    @overload
-    def __abs__(self: NDArray[object_]) -> Any: ...
 
-    @overload
-    def __invert__(self: NDArray[_UnknownType]) -> NDArray[Any]: ...
-    @overload
-    def __invert__(self: NDArray[np.bool]) -> NDArray[np.bool]: ...
-    @overload
-    def __invert__(self: NDArray[_IntType]) -> NDArray[_IntType]: ...
-    @overload
-    def __invert__(self: NDArray[object_]) -> Any: ...
-
-    @overload
-    def __pos__(self: NDArray[_NumberType]) -> NDArray[_NumberType]: ...
-    @overload
-    def __pos__(self: NDArray[timedelta64]) -> NDArray[timedelta64]: ...
-    @overload
-    def __pos__(self: NDArray[object_]) -> Any: ...
-
-    @overload
-    def __neg__(self: NDArray[_NumberType]) -> NDArray[_NumberType]: ...
-    @overload
-    def __neg__(self: NDArray[timedelta64]) -> NDArray[timedelta64]: ...
-    @overload
-    def __neg__(self: NDArray[object_]) -> Any: ...
+    # TODO: Uncomment once https://github.com/python/mypy/issues/14070 is fixed
+    # @overload
+    # def __abs__(self: ndarray[_ShapeType, dtypes.Complex64DType], /) -> ndarray[_ShapeType, dtypes.Float32DType]: ...
+    # @overload
+    # def __abs__(self: ndarray[_ShapeType, dtypes.Complex128DType], /) -> ndarray[_ShapeType, dtypes.Float64DType]: ...
+    # @overload
+    # def __abs__(self: ndarray[_ShapeType, dtypes.CLongDoubleDType], /) -> ndarray[_ShapeType, dtypes.LongDoubleDType]: ...
+    # @overload
+    # def __abs__(self: ndarray[_ShapeType, dtype[complex128]], /) -> ndarray[_ShapeType, dtype[float64]]: ...
+    @overload
+    def __abs__(
+        self: ndarray[_ShapeType, dtype[complexfloating[_NBit_fc]]], /
+    ) -> ndarray[_ShapeType, dtype[floating[_NBit_fc]]]: ...
+    @overload
+    def __abs__(self: _RealArrayT, /) -> _RealArrayT: ...
+    def __invert__(self: _IntegralArrayT, /) -> _IntegralArrayT: ...  # noqa: PYI019
+    def __neg__(self: _NumericArrayT, /) -> _NumericArrayT: ...  # noqa: PYI019
+    def __pos__(self: _NumericArrayT, /) -> _NumericArrayT: ...  # noqa: PYI019
 
     # Binary ops
     @overload
@@ -3094,6 +3070,7 @@ _ScalarType = TypeVar("_ScalarType", bound=generic)
 _NBit = TypeVar("_NBit", bound=NBitBase)
 _NBit1 = TypeVar("_NBit1", bound=NBitBase)
 _NBit2 = TypeVar("_NBit2", bound=NBitBase, default=_NBit1)
+_NBit_fc = TypeVar("_NBit_fc", _NBitHalf, _NBitSingle, _NBitDouble, _NBitLongDouble)
 
 class generic(_ArrayOrScalarCommon):
     @abstractmethod

numpy/typing/tests/data/reveal/ndarray_assignability.pyi

@@ -0,0 +1,79 @@
+from typing import Protocol, TypeAlias, TypeVar
+from typing_extensions import 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)