🎨 DRY cleanup of np.lib._user_array_impl

Author: jorenham

src/numpy-stubs/lib/_user_array_impl.pyi

@@ -1,7 +1,7 @@
 from _typeshed import Incomplete
 from types import EllipsisType
-from typing import Any, Generic, Self, SupportsIndex, TypeAlias, TypeVar, overload
-from typing_extensions import deprecated, override
+from typing import Any, Generic, Self, SupportsIndex, TypeAlias, overload
+from typing_extensions import TypeAliasType, TypeVar, deprecated, override
 
 import _numtype as _nt
 import numpy as np
@@ -10,23 +10,20 @@ from numpy._typing import _ArrayLike, _ArrayLikeBool_co, _ArrayLikeInt_co, _DTyp
 
 ###
 
-_ScalarT = TypeVar("_ScalarT", bound=np.generic)
-_ShapeT = TypeVar("_ShapeT", bound=_nt.Shape)
+_ScalarT = TypeVar("_ScalarT", bound=np.generic, default=Any)
+_ShapeT = TypeVar("_ShapeT", bound=_nt.Shape, default=Any)
 _ShapeT_co = TypeVar("_ShapeT_co", bound=_nt.Shape, default=Any, covariant=True)
-_DTypeT = TypeVar("_DTypeT", bound=np.dtype)
+_DTypeT = TypeVar("_DTypeT", bound=np.dtype, default=np.dtype)
 _DTypeT_co = TypeVar("_DTypeT_co", bound=np.dtype, default=np.dtype, covariant=True)
 
-_BoolArrayT = TypeVar("_BoolArrayT", bound=container[Any, np.dtype[np.bool]])
-_IntegralArrayT = TypeVar("_IntegralArrayT", bound=container[Any, np.dtype[np.bool | np.integer | np.object_]])
-_RealContainerT = TypeVar(
-    "_RealContainerT",
-    bound=container[Any, np.dtype[np.bool | np.integer | np.floating | np.timedelta64 | np.object_]],
-)
-_NumericContainerT = TypeVar(
-    "_NumericContainerT", bound=container[Any, np.dtype[np.number | np.timedelta64 | np.object_]]
-)
+_Container = TypeAliasType("_Container", container[_ShapeT, np.dtype[_ScalarT]], type_params=(_ScalarT, _ShapeT))
 
-_ArrayInt_co: TypeAlias = npt.NDArray[np.integer | np.bool]
+_BoolArrayT = TypeVar("_BoolArrayT", bound=_Container[np.bool_])
+_IntegralArrayT = TypeVar("_IntegralArrayT", bound=_Container[_nt.co_integer | np.object_])
+_RealContainerT = TypeVar("_RealContainerT", bound=_Container[_nt.co_timedelta | np.floating | np.object_])
+_NumericContainerT = TypeVar("_NumericContainerT", bound=_Container[np.number | np.timedelta64 | np.object_])
+
+_ArrayInt_co: TypeAlias = _nt.Array[np.integer | np.bool]
 
 _ToIndexSlice: TypeAlias = slice | EllipsisType | _ArrayInt_co | None
 _ToIndexSlices: TypeAlias = _ToIndexSlice | tuple[_ToIndexSlice, ...]
@@ -48,22 +45,28 @@ class container(Generic[_ShapeT_co, _DTypeT_co]):
     ) -> None: ...
     @overload
     def __init__(
-        self: container[Any, np.dtype[_ScalarT]],
+        self: _Container[_ScalarT],
         /,
         data: _ArrayLike[_ScalarT],
         dtype: None = None,
         copy: bool = True,
     ) -> None: ...
     @overload
     def __init__(
-        self: container[Any, np.dtype[_ScalarT]],
+        self: _Container[_ScalarT],
         /,
         data: npt.ArrayLike,
         dtype: _DTypeLike[_ScalarT],
         copy: bool = True,
     ) -> None: ...
     @overload
-    def __init__(self, /, data: npt.ArrayLike, dtype: npt.DTypeLike | None = None, copy: bool = True) -> None: ...
+    def __init__(
+        self,
+        /,
+        data: npt.ArrayLike,
+        dtype: npt.DTypeLike | None = None,
+        copy: bool = True,
+    ) -> None: ...
 
     #
     def __complex__(self, /) -> complex: ...
@@ -74,15 +77,15 @@ class container(Generic[_ShapeT_co, _DTypeT_co]):
 
     #
     @override
-    def __eq__(self, other: object, /) -> container[_ShapeT_co, np.dtype[np.bool]]: ...  # type: ignore[override]  # pyright: ignore[reportIncompatibleMethodOverride]
+    def __eq__(self, other: object, /) -> _Container[np.bool, _ShapeT_co]: ...  # type: ignore[override]  # pyright: ignore[reportIncompatibleMethodOverride]
     @override
-    def __ne__(self, other: object, /) -> container[_ShapeT_co, np.dtype[np.bool]]: ...  # type: ignore[override]  # pyright: ignore[reportIncompatibleMethodOverride]
+    def __ne__(self, other: object, /) -> _Container[np.bool, _ShapeT_co]: ...  # type: ignore[override]  # pyright: ignore[reportIncompatibleMethodOverride]
 
     #
-    def __lt__(self, other: npt.ArrayLike, /) -> container[_ShapeT_co, np.dtype[np.bool]]: ...
-    def __le__(self, other: npt.ArrayLike, /) -> container[_ShapeT_co, np.dtype[np.bool]]: ...
-    def __gt__(self, other: npt.ArrayLike, /) -> container[_ShapeT_co, np.dtype[np.bool]]: ...
-    def __ge__(self, other: npt.ArrayLike, /) -> container[_ShapeT_co, np.dtype[np.bool]]: ...
+    def __lt__(self, other: npt.ArrayLike, /) -> _Container[np.bool, _ShapeT_co]: ...
+    def __le__(self, other: npt.ArrayLike, /) -> _Container[np.bool, _ShapeT_co]: ...
+    def __gt__(self, other: npt.ArrayLike, /) -> _Container[np.bool, _ShapeT_co]: ...
+    def __ge__(self, other: npt.ArrayLike, /) -> _Container[np.bool, _ShapeT_co]: ...
 
     #
     def __len__(self, /) -> int: ...
@@ -95,27 +98,23 @@ class container(Generic[_ShapeT_co, _DTypeT_co]):
     @overload
     def __getitem__(self, key: _ToIndices, /) -> Any: ...
     @overload
-    def __getitem__(
-        self: container[Any, np.dtype[np.void]], key: list[str], /
-    ) -> container[_ShapeT_co, np.dtype[np.void]]: ...
+    def __getitem__(self: _Container[np.void], key: list[str], /) -> _Container[np.void, _ShapeT_co]: ...
     @overload
-    def __getitem__(self: container[Any, np.dtype[np.void]], key: str, /) -> container[_ShapeT_co, np.dtype]: ...
+    def __getitem__(self: _Container[np.void], key: str, /) -> _Container[Any, _ShapeT_co]: ...
 
     # keep in sync with np.ndarray
     @overload
     def __setitem__(self, index: _ToIndices, value: object, /) -> None: ...
     @overload
-    def __setitem__(self: container[Any, np.dtype[np.void]], key: str | list[str], value: object, /) -> None: ...
+    def __setitem__(self: _Container[np.void], key: str | list[str], value: object, /) -> None: ...
 
     # keep in sync with np.ndarray
     @overload
-    def __abs__(self: container[_ShapeT, np.dtype[np.complex64]], /) -> container[_ShapeT, np.dtype[np.float32]]: ...
+    def __abs__(self: _Container[np.complex64, _ShapeT], /) -> _Container[np.float32, _ShapeT]: ...
     @overload
-    def __abs__(self: container[_ShapeT, np.dtype[np.complex128]], /) -> container[_ShapeT, np.dtype[np.float64]]: ...
+    def __abs__(self: _Container[np.complex128, _ShapeT], /) -> _Container[np.float64, _ShapeT]: ...
     @overload
-    def __abs__(
-        self: container[_ShapeT, np.dtype[np.clongdouble]], /
-    ) -> container[_ShapeT, np.dtype[np.longdouble]]: ...
+    def __abs__(self: _Container[np.clongdouble, _ShapeT], /) -> _Container[np.longdouble, _ShapeT]: ...
     @overload
     def __abs__(self: _RealContainerT, /) -> _RealContainerT: ...
 
@@ -125,7 +124,6 @@ class container(Generic[_ShapeT_co, _DTypeT_co]):
     def __invert__(self: _IntegralArrayT, /) -> _IntegralArrayT: ...  # noqa: PYI019
 
     # TODO(jorenham): complete these binary ops
-
     #
     def __add__(self, other: npt.ArrayLike, /) -> Incomplete: ...
     def __radd__(self, other: npt.ArrayLike, /) -> Incomplete: ...
@@ -161,22 +159,20 @@ class container(Generic[_ShapeT_co, _DTypeT_co]):
     def __ipow__(self, other: npt.ArrayLike, /) -> Self: ...
 
     #
-    def __lshift__(self, other: _ArrayLikeInt_co, /) -> container[Any, np.dtype[np.integer]]: ...
-    def __rlshift__(self, other: _ArrayLikeInt_co, /) -> container[Any, np.dtype[np.integer]]: ...
+    def __lshift__(self, other: _ArrayLikeInt_co, /) -> _Container[np.integer]: ...
+    def __rlshift__(self, other: _ArrayLikeInt_co, /) -> _Container[np.integer]: ...
     def __ilshift__(self, other: _ArrayLikeInt_co, /) -> Self: ...
 
     #
-    def __rshift__(self, other: _ArrayLikeInt_co, /) -> container[Any, np.dtype[np.integer]]: ...
-    def __rrshift__(self, other: _ArrayLikeInt_co, /) -> container[Any, np.dtype[np.integer]]: ...
+    def __rshift__(self, other: _ArrayLikeInt_co, /) -> _Container[np.integer]: ...
+    def __rrshift__(self, other: _ArrayLikeInt_co, /) -> _Container[np.integer]: ...
     def __irshift__(self, other: _ArrayLikeInt_co, /) -> Self: ...
 
     #
     @overload
-    def __and__(
-        self: container[Any, np.dtype[np.bool]], other: _ArrayLikeBool_co, /
-    ) -> container[Any, np.dtype[np.bool]]: ...
+    def __and__(self: _Container[np.bool_], other: _ArrayLikeBool_co, /) -> _Container[np.bool_]: ...
     @overload
-    def __and__(self, other: _ArrayLikeInt_co, /) -> container[Any, np.dtype[np.bool | np.integer]]: ...
+    def __and__(self, other: _ArrayLikeInt_co, /) -> _Container[_nt.co_integer]: ...
     __rand__ = __and__
     @overload
     def __iand__(self: _BoolArrayT, other: _ArrayLikeBool_co, /) -> _BoolArrayT: ...
@@ -185,11 +181,9 @@ class container(Generic[_ShapeT_co, _DTypeT_co]):
 
     #
     @overload
-    def __xor__(
-        self: container[Any, np.dtype[np.bool]], other: _ArrayLikeBool_co, /
-    ) -> container[Any, np.dtype[np.bool]]: ...
+    def __xor__(self: _Container[np.bool_], other: _ArrayLikeBool_co, /) -> _Container[np.bool_]: ...
     @overload
-    def __xor__(self, other: _ArrayLikeInt_co, /) -> container[Any, np.dtype[np.bool | np.integer]]: ...
+    def __xor__(self, other: _ArrayLikeInt_co, /) -> _Container[_nt.co_integer]: ...
     __rxor__ = __xor__
     @overload
     def __ixor__(self: _BoolArrayT, other: _ArrayLikeBool_co, /) -> _BoolArrayT: ...
@@ -198,11 +192,9 @@ class container(Generic[_ShapeT_co, _DTypeT_co]):
 
     #
     @overload
-    def __or__(
-        self: container[Any, np.dtype[np.bool]], other: _ArrayLikeBool_co, /
-    ) -> container[Any, np.dtype[np.bool]]: ...
+    def __or__(self: _Container[np.bool_], other: _ArrayLikeBool_co, /) -> _Container[np.bool_]: ...
     @overload
-    def __or__(self, other: _ArrayLikeInt_co, /) -> container[Any, np.dtype[np.bool | np.integer]]: ...
+    def __or__(self, other: _ArrayLikeInt_co, /) -> _Container[_nt.co_integer]: ...
     __ror__ = __or__
     @overload
     def __ior__(self: _BoolArrayT, other: _ArrayLikeBool_co, /) -> _BoolArrayT: ...
@@ -229,4 +221,4 @@ class container(Generic[_ShapeT_co, _DTypeT_co]):
     def tostring(self, /) -> bytes: ...
     def tobytes(self, /) -> bytes: ...
     def byteswap(self, /) -> Self: ...
-    def astype(self, /, typecode: _DTypeLike[_ScalarT]) -> container[_ShapeT_co, np.dtype[_ScalarT]]: ...
+    def astype(self, /, typecode: _DTypeLike[_ScalarT]) -> _Container[_ScalarT, _ShapeT_co]: ...