-
-
Notifications
You must be signed in to change notification settings - Fork 6
Labels
Description
By doing something like
class int16(signedinteger[_16]):
@type_check_only
def __promote_py__(self) -> builtins.bool: ...
@type_check_only
def __promote_np__(self) -> Self | uint8 | int8 | np.bool: ...
@type_check_only
def __promote_to__(self) -> int16: ... # maybe `Self` also works
...for all generic types (and without overloads), and
class ndarray[ShapeT: Shape, ScalarT: generic]:
@type_check_only
def __promote_py__[T](self: Promotable[T, Any, Any]) -> T: ...
@type_check_only
def __promote_np__[SCT: generic](self: Promotable[Any, SCT, Any]) -> ndarray[ShapeT, dtype[SCT]]: ...
# maybe `Self` also works
@type_check_only
def __promote_to__[RT: generic](self: Promotable[Any, Any, RT]) -> ndarray[ShapeT, dtype[RT]]: ...where
@type_check_only
class HasDType[DT](Protocol): # covariant
@property
def dtype(self) -> DT: ...
@type_check_only
class HasType[T](Protocol): # covariant
@property
def type(self) -> type[T]: ...
@type_check_only
class CanPromote[T, SCT: generic, RT: generic](Protocol): # covariant
def __promote_py__(self) -> T: ...
def __promote_np__(self) -> SCT: ...
def __promote_to__(self) -> RT: ...
type Promotable[T, SCT: generic, RT: Generic] = HasDType[HasType[CanPromote[T, SCT, RT]]]This makes it possible to simply write
@overload
def add[T, SCT: generic, RT: generic](
a: CanPromote[T, SCT, RT],
b: T | SCT,
/,
) -> RT: ...
@overload
def add[T, SCT: generic, RT: generic](
a: T | SCT,
b: CanPromote[T, SCT, RT],
/,
) -> RT: ...
# <remaining (scalar-/array-like, scalar-/array-like) overloads omittedwhere _AsScalarOrArray[T] is a nested protocol for __array__() -> HasDType[HasType[T]]
Without these __promote_*__ methods, the add function would require 4 * (number of scalar types) overloads to achieve the same.
The same trick significantly simplify the ndarray and generic binary operators.