How to type a method decorator that returns an overloaded implementation ?

[QuentinSoubeyranAqemia]

Hi all,

I've been struggling to type the following code and can't understand why MyPy 1.4.0 complains of incompatible signatures.

I'm using the following Protocol, from here, which is reduced to the problematic part is:

import typing
from typing import TypeVar

_T_co = TypeVar("_T_co", covariant=True)

class ArrayLike(Protocol[_T_co]):
    @typing.overload
    def __getitem__(self, index: int, /) -> _T_co:
        ...

    @typing.overload
    def __getitem__(self: Self, index: slice, /) -> "ArrayLike[_T_co]":
        ...

I then need to implement slice support in a lot of place. In most places, the slice support would be simple-minded, and writting the overload annotations, the if...clause and the code becomes very repetitve and tedious. I have written a decorator to do that:

def add_slice_support(getitem):

    def __getitem__(self, index, /):
        if isinstance(index, int):
            return getitem(self, index)
        elif isinstance(index, slice):
            return tuple(getitem(self, i) for i in range(*index.indices(len(self))))
        else:
            raise TypeError()

    return __getitem__

I struggle to write a callback protocol that:

• mypy understands properly (this I managed)
• is compatible with the ArrayLike protocol

I found #1040 but even when adding the descriptor protocol to my callback, MyPy complains it is not compatible with the ArrayLike interface.

This is the full code I have right now:

import typing
from typing import *


_T_co = TypeVar("_T_co", covariant=True)

class ArrayLike(Protocol[_T_co]):
    @typing.overload
    def __getitem__(self, index: int, /) -> _T_co:
        ...

    @typing.overload
    def __getitem__(self, index: slice, /) -> "ArrayLike[_T_co]":
        ...



class _BoundSliceGetitem(Protocol[_T_co]):
    """The bound version of a _SliceGetitem obtained by accessing the method on an instance"""

    @typing.overload
    def __call__(_self_unused, index: int, /) -> _T_co:
        ...

    @typing.overload
    def __call__(_self_unused, index: slice, /) -> ArrayLike[_T_co]:
        ...

    def __call__(_self_unused, index: Union[int, slice], /) -> Union[_T_co, ArrayLike[_T_co]]:
        ...


class _SliceGetitem(Protocol[_T_co]):
    """Protocol for a __getitem__ function that supports slice indexing, including method magic"""

    # Type the magic that happens at python runtime when we type instance.method

    # instance.method access: the `obj` parameter is defined
    @typing.overload
    def __get__(
        self, obj: Sized, objtype: Optional[Type[Sized]] = None
    ) -> _BoundSliceGetitem[_T_co]:
        ...

    # Class.instance access: no instance object, only the class is passed
    @typing.overload
    def __get__(self, obj: None, objtype: Optional[Type[Sized]] = None) -> "_SliceGetitem[_T_co]":
        ...


    def __get__(
        self, obj: Optional[Sized], objtype: Optional[Type[Sized]] = None
    ) -> Union[_BoundSliceGetitem[_T_co], "_SliceGetitem[_T_co]"]:
        ...

    # Type the interface we want. There is an additional `self_` parameter

    @typing.overload
    def __call__(self_, self: Sized, index: int, /) -> _T_co:
        ...

    @typing.overload
    def __call__(self_, self: Sized, index: slice, /) -> Tuple[_T_co, ...]:
        ...


# Now that we have written everything MyPy needs, write the actual implementation of the helper

_Class = TypeVar("_Class", bound=Sized)
_T = TypeVar("_T")


# return type is Any for now because MyPy doesn't think _SliceGetitem is correct
def add_slice_support(getitem: Callable[[_Class, int], _T]) -> _SliceGetitem[_T]:
    """
    Decorator to add support for slice to a __getitem__ supporting only int
    """

    def __getitem__(self: _Class, index: Union[int, slice], /) -> Union[_T, Tuple[_T, ...]]:
        if isinstance(index, int):
            # defer to the underlying implementation
            return getitem(self, index)
        elif isinstance(index, slice):
            # we got a slice, we'll need to support it
            return tuple(getitem(self, i) for i in range(*index.indices(len(self))))
        else:
            raise TypeError("index must be int or slice")

    return typing.cast(_SliceGetitem[_T], __getitem__)


class Test(ArrayLike[int]):
    
    def __len__(self) -> int:
        return 999
    
    # error here: incompatible signature
    @add_slice_support
    def __getitem__(self, index: int, /) -> int:
        return index
        
test = Test()
reveal_type(test[0])
reveal_type(test[90])
reveal_type(test[0:3])
reveal_type(test["invalid"])

MyPy seems to correclty understands the type of the indexing, but still complains at definition site that this is not compatible with ArrayLike.