Skip to content

Embed the promotion rules in the scalar types (!) #137

@jorenham

Description

@jorenham

requires: #136
related: #135


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 omitted

where _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.

Metadata

Metadata

Assignees

Projects

No projects

Milestone

No milestone

Relationships

None yet

Development

No branches or pull requests

Issue actions