✨ HasRank{LE,GE} type aliases for for broadcasting things with a shape

Author: jorenham

src/_numtype/__init__.pyi

@@ -84,8 +84,8 @@ from ._nep50 import (
     CastsWithScalar as CastsWithScalar,
 )
 from ._rank import (
-    Broadcasts as Broadcasts,
-    BroadcastsTo as BroadcastsTo,
+    HasRankGE as HasRankGE,
+    HasRankLE as HasRankLE,
     Rank as Rank,
     Rank0 as Rank0,
     Rank0N as Rank0N,
@@ -97,7 +97,6 @@ from ._rank import (
     Rank3N as Rank3N,
     Rank4 as Rank4,
     Rank4N as Rank4N,
-    _BroadcastableShape as _BroadcastableShape,
 )
 from ._scalar import (
     inexact32 as inexact32,

src/_numtype/_rank.pyi

@@ -4,8 +4,8 @@ from typing_extensions import TypeAliasType, TypeVar
 from ._shape import Shape, Shape0, Shape0N, Shape1, Shape1N, Shape2, Shape2N, Shape3, Shape3N, Shape4, Shape4N
 
 __all__ = [
-    "Broadcasts",
-    "BroadcastsTo",
+    "HasRankGE",
+    "HasRankLE",
     "Rank",
     "Rank0",
     "Rank0N",
@@ -29,22 +29,24 @@ _Shape04: TypeAlias = _Shape03 | Shape4
 
 ###
 
-_ToT = TypeVar("_ToT", bound=Shape)
-_FromT = TypeVar("_FromT", bound=Shape)
+_UpperT = TypeVar("_UpperT", bound=Shape)
+_LowerT = TypeVar("_LowerT", bound=Shape)
 _RankT = TypeVar("_RankT", bound=Shape, default=Any)
 
-_BroadcastableShape = TypeAliasType(
-    "_BroadcastableShape",
-    _FromT | _CanBroadcastFrom[_FromT, _RankT],
-    type_params=(_FromT, _RankT),
+HasRankLE = TypeAliasType(
+    "HasRankLE",
+    _HasShape[_HasOwnShape[_UpperT] | _CanBroadcast[Any, _UpperT, _RankT]],
+    type_params=(_UpperT, _RankT),
+)
+HasRankGE = TypeAliasType(
+    "HasRankGE",
+    _HasShape[_LowerT | _CanBroadcast[_LowerT, Any, _RankT]],
+    type_params=(_LowerT, _RankT),
 )
-
-BroadcastsTo = TypeAliasType("BroadcastsTo", _HasShape[_CanBroadcastTo[_ToT, _RankT]], type_params=(_ToT, _RankT))
-Broadcasts = TypeAliasType("Broadcasts", _HasShape[_BroadcastableShape[_FromT, _RankT]], type_params=(_FromT, _RankT))
 
 ###
 
-_ShapeT_co = TypeVar("_ShapeT_co", bound=Shape | _HasOwnShape | _CanBroadcastFrom | _CanBroadcastTo, covariant=True)
+_ShapeT_co = TypeVar("_ShapeT_co", bound=Shape | _HasOwnShape | _CanBroadcast, covariant=True)
 
 @type_check_only
 class _HasShape(Protocol[_ShapeT_co]):
@@ -53,17 +55,12 @@ class _HasShape(Protocol[_ShapeT_co]):
 
 _FromT_contra = TypeVar("_FromT_contra", default=Any, contravariant=True)
 _ToT_contra = TypeVar("_ToT_contra", bound=Shape, default=Any, contravariant=True)
-_RankT_co = TypeVar("_RankT_co", bound=Shape, default=Any, covariant=True)
-
-@final
-@type_check_only
-class _CanBroadcastFrom(Protocol[_FromT_contra, _RankT_co]):
-    def __broadcast_from__(self, from_: _FromT_contra, /) -> _RankT_co: ...
+_EquivT_co = TypeVar("_EquivT_co", bound=Shape, default=Any, covariant=True)
 
 @final
 @type_check_only
-class _CanBroadcastTo(Protocol[_ToT_contra, _RankT_co]):
-    def __broadcast_to__(self, to: _ToT_contra, /) -> _RankT_co: ...
+class _CanBroadcast(Protocol[_FromT_contra, _ToT_contra, _EquivT_co]):
+    def __broadcast__(self, from_: _FromT_contra, to: _ToT_contra, /) -> _EquivT_co: ...
 
 # 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
@@ -72,6 +69,7 @@ class _CanBroadcastTo(Protocol[_ToT_contra, _RankT_co]):
 
 _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
@@ -81,39 +79,36 @@ class _HasOwnShape(Protocol[_OwnShapeT_contra, _OwnShapeT_co]):
 ###
 # TODO(jorenham): embed the array-like types, e.g. `Sequence[Sequence[T]]`
 
-_OwnShapeT = TypeVar("_OwnShapeT", bound=tuple[Any, ...], default=Any)
-
 @type_check_only
-class _BaseRank(Generic[_FromT_contra, _ToT_contra, _OwnShapeT]):
-    def __broadcast_from__(self, from_: _FromT_contra, /) -> Self: ...
-    def __broadcast_to__(self, to: _ToT_contra, /) -> Self: ...
+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: ...
 
 @type_check_only
 class _BaseRankM(
-    _BaseRank[_FromT_contra | _HasOwnShape[_ToT_contra, Shape], _ToT_contra, _OwnShapeT],
-    Generic[_FromT_contra, _ToT_contra, _OwnShapeT],
+    _BaseRank[_FromT_contra | _HasOwnShape[_ToT_contra, Shape], _OwnShapeT, _ToT_contra],
+    Generic[_FromT_contra, _OwnShapeT, _ToT_contra],
 ): ...
 
 @final
 @type_check_only
-class Rank0(_BaseRankM[_Shape00, Shape0N, Shape0], tuple[()]): ...
+class Rank0(_BaseRankM[_Shape00, Shape0, Shape0N], tuple[()]): ...
 
 @final
 @type_check_only
-class Rank1(_BaseRankM[_Shape01, Shape1N, Shape1], tuple[int]): ...
+class Rank1(_BaseRankM[_Shape01, Shape1, Shape1N], tuple[int]): ...
 
 @final
 @type_check_only
-class Rank2(_BaseRankM[_Shape02, Shape2N, Shape2], tuple[int, int]): ...
+class Rank2(_BaseRankM[_Shape02, Shape2, Shape2N], tuple[int, int]): ...
 
 @final
 @type_check_only
-class Rank3(_BaseRankM[_Shape03, Shape3N, Shape3], tuple[int, int, int]): ...
+class Rank3(_BaseRankM[_Shape03, Shape3, Shape3N], tuple[int, int, int]): ...
 
 @final
 @type_check_only
-class Rank4(_BaseRankM[_Shape04, Shape4N, Shape4], tuple[int, int, int, int]): ...
+class Rank4(_BaseRankM[_Shape04, Shape4, Shape4N], tuple[int, int, int, int]): ...
 
 # this emulates `AnyOf`, rather than a `Union`.
 @type_check_only