✨ gradual rank type

src/_numtype/@test/generated/test_rank.pyi

@@ -1,4 +1,4 @@
-# @generated 2025-05-19T03:25:34Z with tool/testgen.py
+# @generated 2025-09-23T18:34:13Z with tool/testgen.py
 from typing import Any
 
 import _numtype as _nt
@@ -31,7 +31,7 @@ r0n_le_r0n: _nt.HasRankLE[_nt.Rank0N] = r0n
 
 r0_ge_s0: _nt.HasRankGE[_nt.Shape0] = r0
 r0_ge_r0: _nt.HasRankGE[_nt.Rank0] = r0
-r0_ge_s0n: _nt.HasRankGE[_nt.Shape0N] = r0  # type: ignore[assignment]
+r0_ge_s0n: _nt.HasRankGE[_nt.Shape0N] = r0
 r0_ge_r0n: _nt.HasRankGE[_nt.Rank0N] = r0
 r0n_ge_s0: _nt.HasRankGE[_nt.Shape0] = r0n
 r0n_ge_r0: _nt.HasRankGE[_nt.Rank0] = r0n

src/_numtype/@test/test_rank_shape.pyi

@@ -23,3 +23,25 @@ assert_type(a3.shape, _nt.Shape3)  # type: ignore[assert-type]
 a4: _nt.Array4D
 assert_type(a4.__inner_shape__, _nt.Rank4)
 assert_type(a4.shape, _nt.Shape4)  # type: ignore[assert-type]
+
+###
+
+rx: _nt.AnyRank
+r0: _nt.Rank0
+r1: _nt.Rank1
+r2: _nt.Rank2
+r3: _nt.Rank3
+r4: _nt.Rank4
+
+rx0: _nt.AnyRank = r0
+rx1: _nt.AnyRank = r1
+rx2: _nt.AnyRank = r2
+rx3: _nt.AnyRank = r3
+rx4: _nt.AnyRank = r4
+
+# TODO: remove the `# type ignore`s once python/mypy#19908 is fixed (yes, another one)
+
+r0x: _nt.Rank0 = rx  # type: ignore[assignment]
+r1x: _nt.Rank1 = rx  # type: ignore[assignment]
+r2x: _nt.Rank2 = rx  # type: ignore[assignment]
+r3x: _nt.Rank3 = rx  # type: ignore[assignment]

src/_numtype/__init__.pyi

@@ -84,6 +84,7 @@ from ._nep50 import (
     CastsWithScalar as CastsWithScalar,
 )
 from ._rank import (
+    AnyRank as AnyRank,
     HasInnerShape as HasInnerShape,
     HasRankGE as HasRankGE,
     HasRankLE as HasRankLE,

src/_numtype/_rank.pyi

@@ -1,9 +1,10 @@
 from typing import Any, Generic, Protocol, Self, TypeAlias, final, type_check_only
-from typing_extensions import TypeAliasType, TypeVar, TypeVarTuple, override
+from typing_extensions import TypeAliasType, TypeVar, TypeVarTuple
 
 from ._shape import AnyShape, Shape, Shape0, Shape1, Shape1N, Shape2, Shape2N, Shape3, Shape3N, Shape4, Shape4N
 
 __all__ = [
+    "AnyRank",
     "HasInnerShape",
     "HasRankGE",
     "HasRankLE",
@@ -22,11 +23,6 @@ __all__ = [
 
 ###
 
-_Shape01: TypeAlias = Shape0 | Shape1
-_Shape02: TypeAlias = _Shape01 | Shape2
-_Shape03: TypeAlias = _Shape02 | Shape3
-_Shape04: TypeAlias = _Shape03 | Shape4
-
 ###
 
 # TODO(jorenham): remove `| Rank0 | Rank` once python/mypy#19110 is fixed
@@ -42,15 +38,12 @@ _RankGE: TypeAlias = _CanBroadcast[_LowerT, Any, _RankT] | _LowerT | Rank
 HasRankLE = TypeAliasType("HasRankLE", _HasInnerShape[_RankLE[_UpperT, _RankT]], type_params=(_UpperT, _RankT))
 HasRankGE = TypeAliasType("HasRankGE", _HasInnerShape[_RankGE[_LowerT, _RankT]], type_params=(_LowerT, _RankT))
 
-_ShapeT = TypeVar("_ShapeT", bound=Shape)
-
 # for unwrapping potential rank types as shape tuples
+_ShapeT = TypeVar("_ShapeT", bound=Shape)
 HasInnerShape = TypeAliasType("HasInnerShape", _HasInnerShape[_HasOwnShape[Any, _ShapeT]], type_params=(_ShapeT,))
 
 ###
 
-_ShapeLikeT_co = TypeVar("_ShapeLikeT_co", bound=Shape | _HasOwnShape | _CanBroadcast[Any, Any], covariant=True)
-
 _FromT_contra = TypeVar("_FromT_contra", contravariant=True)
 _ToT_contra = TypeVar("_ToT_contra", bound=tuple[Any, ...], contravariant=True)
 _EquivT_co = TypeVar("_EquivT_co", bound=Shape, default=Any, covariant=True)
@@ -61,6 +54,8 @@ _EquivT_co = TypeVar("_EquivT_co", bound=Shape, default=Any, covariant=True)
 class _CanBroadcast(Protocol[_FromT_contra, _ToT_contra, _EquivT_co]):
     def __broadcast__(self, from_: _FromT_contra, to: _ToT_contra, /) -> _EquivT_co: ...
 
+_ShapeLikeT_co = TypeVar("_ShapeLikeT_co", bound=Shape | _HasOwnShape | _CanBroadcast[Any, Any], covariant=True)
+
 # __inner_shape__ is similar to `shape`, but directly exposes the `Rank` type.
 @final
 @type_check_only
@@ -85,72 +80,30 @@ class _HasOwnShape(Protocol[_OwnShapeT_contra, _OwnShapeT_co]):
 
 _Ts = TypeVarTuple("_Ts")  # should only contain `int`s
 
-# https://github.com/python/mypy/issues/19093
-@type_check_only
-class BaseRank(tuple[*_Ts], Generic[*_Ts]):
-    def __broadcast__(self, from_: tuple[*_Ts], to: tuple[*_Ts], /) -> Self: ...
-    def __own_shape__(self, shape: tuple[*_Ts], /) -> tuple[*_Ts]: ...
-
-@final
-@type_check_only
-class Rank0(BaseRank[()]):
-    @override
-    def __broadcast__(self, from_: Shape0 | _HasOwnShape[Shape, Any], to: Shape, /) -> Self: ...
-
-@final
-@type_check_only
-class Rank1(BaseRank[int]):
-    @override
-    def __broadcast__(self, from_: _Shape01 | _HasOwnShape[Shape1N, Any], to: Shape1N, /) -> Self: ...
-
-@final
-@type_check_only
-class Rank2(BaseRank[int, int]):
-    @override
-    def __broadcast__(self, from_: _Shape02 | _HasOwnShape[Shape2N, Any], to: Shape2N, /) -> Self: ...
-
-@final
-@type_check_only
-class Rank3(BaseRank[int, int, int]):
-    @override
-    def __broadcast__(self, from_: _Shape03 | _HasOwnShape[Shape3N, Any], to: Shape3N, /) -> Self: ...
-
 @final
 @type_check_only
-class Rank4(BaseRank[int, int, int, int]):
-    @override
-    def __broadcast__(self, from_: _Shape04 | _HasOwnShape[Shape4N, Any], to: Shape4N, /) -> Self: ...
-
-# these emulates `AnyOf` (gradual union), rather than a `Union`.
-
-@final
-@type_check_only
-class Rank(BaseRank[*tuple[int, ...]]):
-    @override
-    def __broadcast__(self, from_: AnyShape, to: tuple[*_Ts], /) -> Self: ...
+class BaseRank(tuple[*_Ts], Generic[_FromT_contra, _ToT_contra, *_Ts]):
+    def __broadcast__(self, from_: _FromT_contra, to: _ToT_contra, /) -> Self: ...
+    def __own_shape__(self, shape: tuple[*_Ts], /) -> tuple[*_Ts]: ...
 
-@final
-@type_check_only
-class Rank1N(BaseRank[int, *tuple[int, ...]]):
-    @override
-    def __broadcast__(self, from_: AnyShape, to: Shape1N, /) -> Self: ...
+_Shape01: TypeAlias = Shape0 | Shape1
+_Shape02: TypeAlias = _Shape01 | Shape2
+_Shape03: TypeAlias = _Shape02 | Shape3
+_Shape04: TypeAlias = _Shape03 | Shape4
 
-@final
-@type_check_only
-class Rank2N(BaseRank[int, int, *tuple[int, ...]]):
-    @override
-    def __broadcast__(self, from_: AnyShape, to: Shape2N, /) -> Self: ...
+Rank0 = TypeAliasType("Rank0", BaseRank[Shape0 | _HasOwnShape[Shape, Any], Shape, *tuple[()]])
+Rank1 = TypeAliasType("Rank1", BaseRank[_Shape01 | _HasOwnShape[Shape1N, Any], Shape1N, int])
+Rank2 = TypeAliasType("Rank2", BaseRank[_Shape02 | _HasOwnShape[Shape2N, Any], Shape2N, int, int])
+Rank3 = TypeAliasType("Rank3", BaseRank[_Shape03 | _HasOwnShape[Shape3N, Any], Shape3N, int, int, int])
+Rank4 = TypeAliasType("Rank4", BaseRank[_Shape04 | _HasOwnShape[Shape4N, Any], Shape4N, int, int, int, int])
 
-@final
-@type_check_only
-class Rank3N(BaseRank[int, int, int, *tuple[int, ...]]):
-    @override
-    def __broadcast__(self, from_: AnyShape, to: Shape3N, /) -> Self: ...
+Rank = TypeAliasType("Rank", BaseRank[Shape, Shape, *tuple[int, ...]])
+AnyRank = TypeAliasType("AnyRank", BaseRank[Any, AnyShape, *tuple[Any, ...]])
 
-@final
-@type_check_only
-class Rank4N(BaseRank[int, int, int, int, *tuple[int, ...]]):
-    @override
-    def __broadcast__(self, from_: AnyShape, to: Shape4N, /) -> Self: ...
+Rank0N = Rank
+Rank1N = TypeAliasType("Rank1N", BaseRank[AnyShape, Shape1N, int, *tuple[int, ...]])
+Rank2N = TypeAliasType("Rank2N", BaseRank[AnyShape, Shape2N, int, int, *tuple[int, ...]])
+Rank3N = TypeAliasType("Rank3N", BaseRank[AnyShape, Shape3N, int, int, int, *tuple[int, ...]])
+Rank4N = TypeAliasType("Rank4N", BaseRank[AnyShape, Shape4N, int, int, int, int, *tuple[int, ...]])
 
-Rank0N: TypeAlias = Rank
+# these emulate `AnyOf` (gradual union), rather than a `Union`.

src/numpy-stubs/@test/static/accept/fromnumeric.pyi

@@ -144,8 +144,7 @@ assert_type(np.nonzero(AR_f4), tuple[_nt.Array[np.intp], ...])
 assert_type(np.nonzero(AR_1d), tuple[_nt.Array[np.intp], ...])
 assert_type(np.nonzero(AR_nd), tuple[_nt.Array[np.intp], ...])
 
-# TODO: remove the `# type: ignore` once python/mypy#19110 is fixed
-assert_type(np.shape(b1), tuple[()])  # type: ignore[assert-type]
+assert_type(np.shape(b1), tuple[()])
 assert_type(np.shape(f_0d), tuple[()])
 assert_type(np.shape(i_1d), tuple[int])
 assert_type(np.shape(i_2d), tuple[int, int])

src/numpy-stubs/__init__.pyi

@@ -2711,7 +2711,7 @@ class ndarray(_ArrayOrScalarCommon, Generic[_ShapeT_co, _DTypeT_co]):
     #
     @override
     @overload  # workaround for python/mypy#19110
-    def tolist(self: _HasShapeAndItem[_nt.Rank0, _T], /) -> _T: ...  # type: ignore[overload-overlap]
+    def tolist(self: _HasShapeAndItem[_nt.Rank0, _T], /) -> _T: ...
     @overload  # workaround for microsoft/pyright#10232
     def tolist(self: ndarray[_nt.NeitherShape], /) -> Any: ...
     @overload

src/numpy-stubs/_core/_multiarray_umath.pyi

@@ -641,7 +641,7 @@ def empty_like(
     device: _Device | None = None,
 ) -> _nt.Array[np.bool_]: ...
 @overload  # bool 0d array-like
-def empty_like(  # type: ignore[overload-overlap]
+def empty_like(
     prototype: _nt.ToBool_0d,
     dtype: _nt.ToDTypeBool | None = None,
     order: _OrderKACF = "K",
@@ -681,7 +681,7 @@ def empty_like(
     device: _Device | None = None,
 ) -> _nt.Array3D[np.bool_]: ...
 @overload  # workaround for microsoft/pyright#10232
-def empty_like(
+def empty_like(  # type: ignore[overload-overlap]  # python/mypy#19908
     prototype: _nt._ToArray_nnd[np.intp],
     dtype: _nt.ToDTypeInt64 | None = None,
     order: _OrderKACF = "K",
@@ -691,7 +691,7 @@ def empty_like(
     device: _Device | None = None,
 ) -> _nt.Array[np.intp]: ...
 @overload  # int 0d array-like
-def empty_like(  # type: ignore[overload-overlap]
+def empty_like(
     prototype: _nt.ToInt_0d,
     dtype: _nt.ToDTypeInt64 | None = None,
     order: _OrderKACF = "K",
@@ -731,7 +731,7 @@ def empty_like(
     device: _Device | None = None,
 ) -> _nt.Array3D[np.intp]: ...
 @overload  # workaround for microsoft/pyright#10232
-def empty_like(
+def empty_like(  # type: ignore[overload-overlap]  # python/mypy#19908
     prototype: _nt._ToArray_nnd[np.float64],
     dtype: _nt.ToDTypeFloat64 | None = None,
     order: _OrderKACF = "K",
@@ -741,7 +741,7 @@ def empty_like(
     device: _Device | None = None,
 ) -> _nt.Array[np.float64]: ...
 @overload  # float 0d array-like
-def empty_like(  # type: ignore[overload-overlap]
+def empty_like(
     prototype: _nt.ToFloat64_0d,
     dtype: _nt.ToDTypeFloat64 | None = None,
     order: _OrderKACF = "K",
@@ -781,7 +781,7 @@ def empty_like(
     device: _Device | None = None,
 ) -> _nt.Array3D[np.float64]: ...
 @overload  # complex 0d array-like
-def empty_like(  # type: ignore[overload-overlap]
+def empty_like(
     prototype: _nt.ToComplex128_0d,
     dtype: _nt.ToDTypeComplex128 | None = None,
     order: _OrderKACF = "K",

src/numpy-stubs/_core/fromnumeric.pyi

@@ -422,8 +422,6 @@ def squeeze(a: _ScalarT, axis: _ShapeLike | None = None) -> _nt.Array0D[_ScalarT
 @overload  # workaround for microsoft/pyright#10232
 def squeeze(a: _nt._ToArray_nnd[_ScalarT], axis: _ShapeLike | None = None) -> _nt.Array[_ScalarT]: ...
 @overload
-def squeeze(a: _nt.CanArray0D[_ScalarT], axis: _ShapeLike | None = None) -> _nt.Array0D[_ScalarT]: ...
-@overload
 def squeeze(a: _nt._ToArray_nd[_ScalarT], axis: _ShapeLike | None = None) -> _nt.Array[_ScalarT]: ...
 @overload
 def squeeze(a: ArrayLike, axis: _ShapeLike | None = None) -> _nt.Array[Incomplete]: ...