🎨 refactor from _numtype import in lib

Author: jorenham

src/numpy-stubs/lib/_arraysetops_impl.pyi

@@ -1,26 +1,10 @@
 from _typeshed import Incomplete
 from collections.abc import Sequence
-from typing import Any, ClassVar, Final, Generic, Literal as L, SupportsIndex, TypeAlias, final, overload
+from typing import Any, ClassVar, Final, Generic, Literal as L, SupportsIndex as CanIndex, TypeAlias, final, overload
 from typing_extensions import Self, TypeVar, deprecated
 
+import _numtype as _nt
 import numpy as np
-from _numtype import (
-    Array,
-    JustBytes,
-    JustComplex,
-    JustFloat,
-    JustInt,
-    JustStr,
-    Sequence1ND,
-    ToBool_nd,
-    ToBytes_nd,
-    ToComplex128_nd,
-    ToFloat64_nd,
-    ToInt_nd,
-    ToObject_nd,
-    ToStr_nd,
-    _ToArray_nd,
-)
 from numpy._core.multiarray import ravel_multi_index, unravel_index
 from numpy._typing import ArrayLike, _SupportsDType as _HasDType
 
@@ -46,7 +30,7 @@ _DTypeT = TypeVar("_DTypeT", bound=np.dtype)
 _ScalarT = TypeVar("_ScalarT", bound=np.generic)
 _ScalarT_co = TypeVar("_ScalarT_co", bound=np.generic, default=Any, covariant=True)
 _TupleT = TypeVar("_TupleT", bound=tuple[object, ...])
-_ArrayT = TypeVar("_ArrayT", bound=Array)
+_ArrayT = TypeVar("_ArrayT", bound=_nt.Array)
 
 _BoolT_co = TypeVar("_BoolT_co", bound=bool, default=bool, covariant=True)
 
@@ -55,32 +39,32 @@ _MatrixT_co = TypeVar("_MatrixT_co", bound=bool, default=L[False], covariant=Tru
 _NDMinT_co = TypeVar("_NDMinT_co", bound=int, default=L[1], covariant=True)
 _Trans1DT_co = TypeVar("_Trans1DT_co", bound=int, default=L[-1], covariant=True)
 
-_Arrays: TypeAlias = tuple[Array[_ScalarT], ...]
+_Arrays: TypeAlias = tuple[_nt.Array[_ScalarT], ...]
 
 ###
 
 class ndenumerate(Generic[_ScalarT_co]):
     @overload
-    def __init__(self: ndenumerate[_ScalarT], /, arr: _ToArray_nd[_ScalarT]) -> None: ...
+    def __init__(self: ndenumerate[_ScalarT], /, arr: _nt._ToArray_nd[_ScalarT]) -> None: ...
     @overload
-    def __init__(self: ndenumerate[np.bytes_], /, arr: ToBytes_nd) -> None: ...
+    def __init__(self: ndenumerate[np.bytes_], /, arr: _nt.ToBytes_nd) -> None: ...
     @overload
-    def __init__(self: ndenumerate[np.str_], /, arr: ToStr_nd) -> None: ...
+    def __init__(self: ndenumerate[np.str_], /, arr: _nt.ToStr_nd) -> None: ...
     @overload
-    def __init__(self: ndenumerate[np.bool], /, arr: ToBool_nd) -> None: ...
+    def __init__(self: ndenumerate[np.bool], /, arr: _nt.ToBool_nd) -> None: ...
     @overload
-    def __init__(self: ndenumerate[np.intp], /, arr: ToInt_nd) -> None: ...
+    def __init__(self: ndenumerate[np.intp], /, arr: _nt.ToInt_nd) -> None: ...
     @overload
-    def __init__(self: ndenumerate[np.float64], /, arr: ToFloat64_nd) -> None: ...
+    def __init__(self: ndenumerate[np.float64], /, arr: _nt.ToFloat64_nd) -> None: ...
     @overload
-    def __init__(self: ndenumerate[np.complex128], /, arr: ToComplex128_nd) -> None: ...
+    def __init__(self: ndenumerate[np.complex128], /, arr: _nt.ToComplex128_nd) -> None: ...
     @overload
-    def __init__(self: ndenumerate[np.object_], /, arr: ToObject_nd) -> None: ...
+    def __init__(self: ndenumerate[np.object_], /, arr: _nt.ToObject_nd) -> None: ...
 
     # The first overload is a (semi-)workaround for a mypy bug (tested with v1.10 and v1.11)
     @overload
     def __next__(
-        self: ndenumerate[np.bool | np.number | np.flexible | np.datetime64 | np.timedelta64],
+        self: ndenumerate[_nt.co_number | np.flexible | np.datetime64 | np.timedelta64],
         /,
     ) -> tuple[tuple[int, ...], _ScalarT_co]: ...
     @overload
@@ -93,9 +77,9 @@ class ndenumerate(Generic[_ScalarT_co]):
 
 class ndindex:
     @overload
-    def __init__(self, shape: tuple[SupportsIndex, ...], /) -> None: ...
+    def __init__(self, shape: tuple[CanIndex, ...], /) -> None: ...
     @overload
-    def __init__(self, /, *shape: SupportsIndex) -> None: ...
+    def __init__(self, /, *shape: CanIndex) -> None: ...
 
     #
     def __iter__(self) -> Self: ...
@@ -107,9 +91,9 @@ class nd_grid(Generic[_BoolT_co]):
     sparse: _BoolT_co
     def __init__(self, sparse: _BoolT_co = ...) -> None: ...
     @overload
-    def __getitem__(self: nd_grid[L[False]], key: slice | Sequence[slice]) -> Array: ...
+    def __getitem__(self: nd_grid[L[False]], key: slice | Sequence[slice]) -> _nt.Array: ...
     @overload
-    def __getitem__(self: nd_grid[L[True]], key: slice | Sequence[slice]) -> tuple[Array, ...]: ...
+    def __getitem__(self: nd_grid[L[True]], key: slice | Sequence[slice]) -> tuple[_nt.Array, ...]: ...
 
 @final
 class MGridClass(nd_grid[L[False]]):
@@ -146,13 +130,17 @@ class AxisConcatenator(Generic[_AxisT_co, _MatrixT_co, _NDMinT_co, _Trans1DT_co]
     #
     @staticmethod
     @overload
-    def concatenate(*a: _ToArray_nd[_ScalarT], axis: SupportsIndex | None = 0, out: None = None) -> Array[_ScalarT]: ...
+    def concatenate(
+        *a: _nt._ToArray_nd[_ScalarT],
+        axis: CanIndex | None = 0,
+        out: None = None,
+    ) -> _nt.Array[_ScalarT]: ...
     @staticmethod
     @overload
-    def concatenate(*a: ArrayLike, axis: SupportsIndex | None = 0, out: _ArrayT) -> _ArrayT: ...
+    def concatenate(*a: ArrayLike, axis: CanIndex | None = 0, out: _ArrayT) -> _ArrayT: ...
     @staticmethod
     @overload
-    def concatenate(*a: ArrayLike, axis: SupportsIndex | None = 0, out: None = None) -> Array: ...
+    def concatenate(*a: ArrayLike, axis: CanIndex | None = 0, out: None = None) -> _nt.Array: ...
 
 @final
 class RClass(AxisConcatenator[L[0], L[False], L[1], L[-1]]):
@@ -176,31 +164,23 @@ class IndexExpression(Generic[_BoolT_co]):
     def __getitem__(self, item: _T, /) -> _T | tuple[_T]: ...
 
 @overload
-def ix_(
-    *args: _HasDType[_DTypeT] | Sequence1ND[_HasDType[_DTypeT]],
-) -> tuple[np.ndarray[tuple[int, ...], _DTypeT], ...]: ...
+def ix_(*args: _nt.SequenceND[_HasDType[_DTypeT]]) -> tuple[np.ndarray[tuple[int, ...], _DTypeT], ...]: ...
 @overload
-def ix_(*args: bool | Sequence1ND[bool]) -> _Arrays[np.bool]: ...
+def ix_(*args: _nt.SequenceND[bool]) -> _Arrays[np.bool]: ...
 @overload
-def ix_(*args: JustInt | Sequence1ND[JustInt]) -> _Arrays[np.intp]: ...
+def ix_(*args: _nt.SequenceND[_nt.JustInt]) -> _Arrays[np.intp]: ...
 @overload
-def ix_(*args: JustFloat | Sequence1ND[JustFloat]) -> _Arrays[np.float64]: ...
+def ix_(*args: _nt.SequenceND[_nt.JustFloat]) -> _Arrays[np.float64]: ...
 @overload
-def ix_(*args: JustComplex | Sequence1ND[JustComplex]) -> _Arrays[np.complex128]: ...
+def ix_(*args: _nt.SequenceND[_nt.JustComplex]) -> _Arrays[np.complex128]: ...
 @overload
-def ix_(*args: JustBytes | Sequence1ND[JustBytes]) -> _Arrays[np.bytes_]: ...
+def ix_(*args: _nt.SequenceND[_nt.JustBytes]) -> _Arrays[np.bytes_]: ...
 @overload
-def ix_(*args: JustStr | Sequence1ND[JustStr]) -> _Arrays[np.str_]: ...
-@overload
-def ix_(*args: int | Sequence1ND[int]) -> _Arrays[np.intp | np.bool]: ...  # type: ignore[overload-cannot-match]  # pyright: ignore[reportOverlappingOverload]
-@overload
-def ix_(*args: float | Sequence1ND[float]) -> _Arrays[np.float64 | np.intp | np.bool]: ...  # type: ignore[overload-cannot-match]  # pyright: ignore[reportOverlappingOverload]
-@overload
-def ix_(*args: complex | Sequence1ND[complex]) -> _Arrays[np.complex128 | np.float64 | np.intp | np.bool]: ...  # type: ignore[overload-cannot-match]  # pyright: ignore[reportOverlappingOverload]
+def ix_(*args: _nt.SequenceND[_nt.JustStr]) -> _Arrays[np.str_]: ...
 
 #
-def fill_diagonal(a: Array, val: Any, wrap: bool = ...) -> None: ...
-def diag_indices(n: int, ndim: int = ...) -> _Arrays[np.intp]: ...
+def fill_diagonal(a: _nt.Array, val: Any, wrap: bool = False) -> None: ...
+def diag_indices(n: int, ndim: int = 2) -> _Arrays[np.intp]: ...
 def diag_indices_from(arr: ArrayLike) -> _Arrays[np.intp]: ...
 
 ###