♻️ prefer associated-ish over generic types rank types

Author: jorenham

src/_numtype/@test/test_rank_shape.pyi

@@ -0,0 +1,28 @@
+from typing import assert_type
+
+import _numtype as _nt
+
+a0: _nt.Array0D
+assert_type(a0.__inner_shape__, _nt.Rank0)
+r0: _nt.Rank0 = a0.shape  # type: ignore[assignment]  # pyright: ignore[reportAssignmentType]
+s0: _nt.Shape0 = a0.shape
+
+a1: _nt.Array1D
+assert_type(a1.__inner_shape__, _nt.Rank1)
+r1: _nt.Rank1 = a1.shape  # type: ignore[assignment]  # pyright: ignore[reportAssignmentType]
+s1: _nt.Shape1 = a1.shape
+
+a2: _nt.Array2D
+assert_type(a2.__inner_shape__, _nt.Rank2)
+r2: _nt.Rank2 = a2.shape  # type: ignore[assignment]  # pyright: ignore[reportAssignmentType]
+s2: _nt.Shape2 = a2.shape
+
+a3: _nt.Array3D
+assert_type(a3.__inner_shape__, _nt.Rank3)
+r3: _nt.Rank3 = a3.shape  # type: ignore[assignment]  # pyright: ignore[reportAssignmentType]
+s3: _nt.Shape3 = a3.shape
+
+a4: _nt.Array4D
+assert_type(a4.__inner_shape__, _nt.Rank4)
+r4: _nt.Rank4 = a4.shape  # type: ignore[assignment]  # pyright: ignore[reportAssignmentType]
+s4: _nt.Shape4 = a4.shape

src/_numtype/__init__.pyi

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

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
+from typing_extensions import TypeAliasType, TypeVar, TypeVarTuple, override
 
 from ._shape import Shape, Shape0, Shape0N, Shape1, Shape1N, Shape2, Shape2N, Shape3, Shape3N, Shape4, Shape4N
 
 __all__ = [
+    "HasInnerShape",
     "HasRankGE",
     "HasRankLE",
     "Rank",
@@ -21,8 +22,7 @@ __all__ = [
 
 ###
 
-_Shape00: TypeAlias = Shape0
-_Shape01: TypeAlias = _Shape00 | Shape1
+_Shape01: TypeAlias = Shape0 | Shape1
 _Shape02: TypeAlias = _Shape01 | Shape2
 _Shape03: TypeAlias = _Shape02 | Shape3
 _Shape04: TypeAlias = _Shape03 | Shape4
@@ -33,44 +33,57 @@ _UpperT = TypeVar("_UpperT", bound=Shape)
 _LowerT = TypeVar("_LowerT", bound=Shape)
 _RankT = TypeVar("_RankT", bound=Shape, default=Any)
 
+_RankLE: TypeAlias = _CanBroadcast[Any, _UpperT, _RankT]
+_RankGE: TypeAlias = _CanBroadcast[_LowerT, Any, _RankT]
+
 HasRankLE = TypeAliasType(
     "HasRankLE",
-    _HasShape[Shape0 | _HasOwnShape[_UpperT] | _CanBroadcast[Any, _UpperT, _RankT]],
+    _HasInnerShape[Shape0 | _RankLE[_UpperT, _RankT]],
     type_params=(_UpperT, _RankT),
 )
 HasRankGE = TypeAliasType(
     "HasRankGE",
-    _HasShape[_LowerT | _CanBroadcast[_LowerT, Any, _RankT]],
+    _HasInnerShape[_LowerT | _RankGE[_LowerT, _RankT]],
     type_params=(_LowerT, _RankT),
 )
 
-###
+_ShapeT = TypeVar("_ShapeT", bound=Shape)
 
-_ShapeT_co = TypeVar("_ShapeT_co", bound=Shape | _HasOwnShape | _CanBroadcast, covariant=True)
+# for unwrapping potential rank types as shape tuples
+HasInnerShape = TypeAliasType(
+    "HasInnerShape",
+    _HasInnerShape[_HasOwnShape[Any, _ShapeT]],
+    type_params=(_ShapeT,),
+)
 
-@type_check_only
-class _HasShape(Protocol[_ShapeT_co]):
-    @property
-    def shape(self, /) -> _ShapeT_co: ...
+###
+
+_ShapeLikeT_co = TypeVar("_ShapeLikeT_co", bound=Shape | _HasOwnShape | _CanBroadcast[Any, Any], covariant=True)
 
-_FromT_contra = TypeVar("_FromT_contra", default=Any, contravariant=True)
-_ToT_contra = TypeVar("_ToT_contra", bound=Shape, default=Any, contravariant=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)
 
+# __broadcast__ is the type-check-only interface order of ranks
 @final
 @type_check_only
 class _CanBroadcast(Protocol[_FromT_contra, _ToT_contra, _EquivT_co]):
     def __broadcast__(self, from_: _FromT_contra, to: _ToT_contra, /) -> _EquivT_co: ...
 
+# __inner_shape__ is similar to `shape`, but directly exposes the `Rank` type.
+@final
+@type_check_only
+class _HasInnerShape(Protocol[_ShapeLikeT_co]):
+    @property
+    def __inner_shape__(self, /) -> _ShapeLikeT_co: ...
+
+_OwnShapeT_contra = TypeVar("_OwnShapeT_contra", bound=tuple[Any, ...], default=Any, contravariant=True)
+_OwnShapeT_co = TypeVar("_OwnShapeT_co", bound=Shape, default=_OwnShapeT_contra, covariant=True)
+
 # This double shape-type parameter is a sneaky way to annotate a doubly-bound nominal type range,
 # e.g. `_HasOwnShape[Shape2N, Shape0N]` accepts `Shape2N`, `Shape1N`, and `Shape0N`, but
 # rejects `Shape3N` and `Shape1`. Besides brevity, it also works around several mypy bugs that
 # are related to "unions vs joins".
-
-_OwnShapeT_contra = TypeVar("_OwnShapeT_contra", bound=Shape, default=Any, contravariant=True)
-_OwnShapeT_co = TypeVar("_OwnShapeT_co", bound=Shape, default=_OwnShapeT_contra, covariant=True)
-_OwnShapeT = TypeVar("_OwnShapeT", bound=tuple[Any, ...], default=Any)
-
 @final
 @type_check_only
 class _HasOwnShape(Protocol[_OwnShapeT_contra, _OwnShapeT_co]):
@@ -79,59 +92,74 @@ class _HasOwnShape(Protocol[_OwnShapeT_contra, _OwnShapeT_co]):
 ###
 # TODO(jorenham): embed the array-like types, e.g. `Sequence[Sequence[T]]`
 
-@type_check_only
-class _BaseRank(Generic[_FromT_contra, _OwnShapeT, _ToT_contra]):
-    def __broadcast__(self, from_: _FromT_contra, to: _ToT_contra, /) -> Self: ...
-    def __own_shape__(self, shape: _OwnShapeT, /) -> _OwnShapeT: ...
+_Ts = TypeVarTuple("_Ts")  # should only contain `int`s
 
+# https://github.com/python/mypy/issues/19093
 @type_check_only
-class _BaseRankM(
-    _BaseRank[_FromT_contra | _HasOwnShape[_ToT_contra, Shape], _OwnShapeT, _ToT_contra],
-    Generic[_FromT_contra, _OwnShapeT, _ToT_contra],
-): ...
+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(_BaseRankM[_Shape00, Shape0, Shape0N], tuple[()]): ...
+class Rank0(BaseRank[*tuple[()]]):
+    @override
+    def __broadcast__(self, from_: Shape0 | _HasOwnShape[Shape, Any], to: Shape, /) -> Self: ...
 
 @final
 @type_check_only
-class Rank1(_BaseRankM[_Shape01, Shape1, Shape1N], tuple[int]): ...
+class Rank1(BaseRank[int]):
+    @override
+    def __broadcast__(self, from_: _Shape01 | _HasOwnShape[Shape1N, Any], to: Shape1N, /) -> Self: ...
 
 @final
 @type_check_only
-class Rank2(_BaseRankM[_Shape02, Shape2, Shape2N], tuple[int, int]): ...
+class Rank2(BaseRank[int, int]):
+    @override
+    def __broadcast__(self, from_: _Shape02 | _HasOwnShape[Shape2N, Any], to: Shape2N, /) -> Self: ...
 
 @final
 @type_check_only
-class Rank3(_BaseRankM[_Shape03, Shape3, Shape3N], tuple[int, int, int]): ...
+class Rank3(BaseRank[int, int, int]):
+    @override
+    def __broadcast__(self, from_: _Shape03 | _HasOwnShape[Shape3N, Any], to: Shape3N, /) -> Self: ...
 
 @final
 @type_check_only
-class Rank4(_BaseRankM[_Shape04, Shape4, Shape4N], tuple[int, int, int, int]): ...
+class Rank4(BaseRank[int, int, int, int]):
+    @override
+    def __broadcast__(self, from_: _Shape04 | _HasOwnShape[Shape4N, Any], to: Shape4N, /) -> Self: ...
 
-# this emulates `AnyOf`, rather than a `Union`.
-@type_check_only
-class _BaseRankMToN(_BaseRank[Shape0N, _OwnShapeT, _OwnShapeT], Generic[_OwnShapeT]): ...
+# these emulates `AnyOf` (gradual union), rather than a `Union`.
 
 @final
 @type_check_only
-class Rank(_BaseRankMToN[Shape0N], tuple[int, ...]): ...
+class Rank(BaseRank[*tuple[int, ...]]):
+    @override
+    def __broadcast__(self, from_: Shape0N, to: tuple[*_Ts], /) -> Self: ...
 
 @final
 @type_check_only
-class Rank1N(_BaseRankMToN[Shape1N], tuple[int, *tuple[int, ...]]): ...
+class Rank1N(BaseRank[int, *tuple[int, ...]]):
+    @override
+    def __broadcast__(self, from_: Shape0N, to: Shape1N, /) -> Self: ...
 
 @final
 @type_check_only
-class Rank2N(_BaseRankMToN[Shape2N], tuple[int, int, *tuple[int, ...]]): ...
+class Rank2N(BaseRank[int, int, *tuple[int, ...]]):
+    @override
+    def __broadcast__(self, from_: Shape0N, to: Shape2N, /) -> Self: ...
 
 @final
 @type_check_only
-class Rank3N(_BaseRankMToN[Shape3N], tuple[int, int, int, *tuple[int, ...]]): ...
+class Rank3N(BaseRank[int, int, int, *tuple[int, ...]]):
+    @override
+    def __broadcast__(self, from_: Shape0N, to: Shape3N, /) -> Self: ...
 
 @final
 @type_check_only
-class Rank4N(_BaseRankMToN[Shape4N], tuple[int, int, int, int, *tuple[int, ...]]): ...
+class Rank4N(BaseRank[int, int, int, int, *tuple[int, ...]]):
+    @override
+    def __broadcast__(self, from_: Shape0N, to: Shape4N, /) -> Self: ...
 
 Rank0N: TypeAlias = Rank