Type-Checkers fail for overloads based on Literal attribute.

[randolf-scholz]

Type hinting this code has been driving me nuts, I can't seem to get it to parse with either pyright or mypy, it could even be that there are some bugs.

I have a class that, depending on its state (attribute "mode"), returns different things when calling __iter__. To represent this, I am using a constrained TypeVar.

Issues with mypy: It seems that mypy selects the wrong overload when determining the type of iter(sampler). In particular, iter(sampler) and sampler.__iter__() reveal different types!! I think this may be a bug and have opened python/mypy#16552

Issues with pyright: With the code as it is below, pyright gives the wrong result for the unknown case (if the mode is not known statically it should fall back to Any). This is expected since pyright selects the first matching overload. Now to circumvent this I have considered adding an extra overload def __iter__(self: "Sliding[T, Never]") -> Iterator[Any]: ..., but when I do that, then the slice case breaks, and curiously only the slice case, irrespective of the overload order of the modes. edit: pyright selects the first option from the Modes type variable in this case. (microsoft/pyright#6557)

Full code example

from typing import (
    TYPE_CHECKING,
    Any,
    Generic,
    Iterable,
    Iterator,
    Literal,
    Never,
    TypeAlias,
    TypeVar,
    assert_type,
    overload,
)

T = TypeVar("T")
S: TypeAlias = Literal["slices"]
B: TypeAlias = Literal["bounds"]
M: TypeAlias = Literal["masks"]
W: TypeAlias = Literal["windows"]
Mode = TypeVar("Mode", S, B, M, W)


class Sliding(Generic[T, Mode]):
    data: list[T]
    size: int
    mode: Mode

    @overload
    def __init__(self: "Sliding[T, S]", data: list[T], size: int, mode: S) -> None: ...
    @overload
    def __init__(self: "Sliding[T, B]", data: list[T], size: int, mode: B) -> None: ...
    @overload
    def __init__(self: "Sliding[T, M]", data: list[T], size: int, mode: M) -> None: ...
    @overload
    def __init__(self: "Sliding[T, W]", data: list[T], size: int, mode: W) -> None: ...
    @overload  # fallback
    def __init__(
        self: "Sliding[T, Any]", data: list[T], size: int, mode: str
    ) -> None: ...
    def __init__(self, data, size, mode):
        self.data = data
        self.size = size
        self.mode = mode

    # @overload  # early fallback
    # def __iter__(self: "Sliding[T, Never]") -> Iterator[Any]: ...
    @overload
    def __iter__(self: "Sliding[T, B]") -> Iterator[tuple[T, T]]: ...
    @overload
    def __iter__(self: "Sliding[T, W]") -> Iterator[list[T]]: ...
    @overload
    def __iter__(self: "Sliding[T, M]") -> Iterator[list[bool]]: ...
    @overload
    def __iter__(self: "Sliding[T, S]") -> Iterator[slice]: ...
    @overload  # late fallback
    def __iter__(self: "Sliding[T, Any]") -> Iterator[Any]: ...
    def __iter__(self):  # pyright: ignore[reportGeneralTypeIssues]
        num = len(self.data) - self.size + 1

        match self.mode:
            case "slices":
                for i in range(num):
                    yield slice(i, i + self.size)
            case "windows":
                for i in range(num):
                    yield self.data[i : i + self.size]
            case "bounds":
                for i in range(num):
                    yield i, i + self.size
            case "masks":
                for i in range(num):
                    yield [i <= j < i + self.size for j in range(len(self.data))]
            case _:
                raise TypeError(f"Unknown mode: {self.mode}")


def test_slice() -> None:
    sampler = Sliding([1, 2, 3, 4, 5], 3, mode="slices")

    if TYPE_CHECKING:
        assert_type(sampler, Sliding[int, S])
        assert_type(sampler.__iter__(), Iterator[slice])
        assert_type(iter(sampler), Iterator[slice])
        assert_type(list(sampler), list[slice])
        s1: Iterable[slice] = sampler

    assert list(sampler) == [slice(0, 3), slice(1, 4), slice(2, 5)]


def test_window() -> None:
    sampler = Sliding([1, 2, 3, 4, 5], 3, mode="windows")

    if TYPE_CHECKING:
        assert_type(sampler, Sliding[int, W])
        assert_type(sampler.__iter__(), Iterator[list[int]])
        assert_type(iter(sampler), Iterator[list[int]])
        assert_type(list(sampler), list[list[int]])
        w1: Iterable[list[int]] = sampler

    assert list(sampler) == [[1, 2, 3], [2, 3, 4], [3, 4, 5]]


def test_bounds() -> None:
    sampler = Sliding([1, 2, 3, 4, 5], 3, mode="bounds")

    if TYPE_CHECKING:
        assert_type(sampler, Sliding[int, B])
        assert_type(sampler.__iter__(), Iterator[tuple[int, int]])
        assert_type(iter(sampler), Iterator[tuple[int, int]])
        assert_type(list(sampler), list[tuple[int, int]])
        b1: Iterable[tuple[int, int]] = sampler

    assert list(sampler) == [(0, 3), (1, 4), (2, 5)]


def test_masks() -> None:
    sampler = Sliding([1, 2, 3, 4, 5], 3, mode="masks")

    if TYPE_CHECKING:
        assert_type(sampler, Sliding[int, M])
        assert_type(sampler.__iter__(), Iterator[list[bool]])
        assert_type(iter(sampler), Iterator[list[bool]])
        assert_type(list(sampler), list[list[bool]])
        m1: Iterable[list[bool]] = sampler

    assert list(sampler) == [
        [True, True, True, False, False],
        [False, True, True, True, False],
        [False, False, True, True, True],
    ]


def test_unknown() -> None:
    mode: str = "..."
    sampler = Sliding([1, 2, 3, 4, 5], 3, mode=mode)

    if TYPE_CHECKING:
        assert_type(sampler, Sliding[int, Any])
        assert_type(sampler.__iter__(), Iterator[Any])
        assert_type(iter(sampler), Iterator[Any])
        assert_type(list(sampler), list[Any])
        u1: Iterable[Any] = sampler

[Daverball]

The iter list thing looks like a bug to me as well, as for the unknown case, it looks like the behavior with using S | B | W | M instead of Any as Mode for the unknown case helps and will return the correct overload for __iter__. I think that's a better representation of the unknown mode anyways, and still allows you to return Iterator[Any].

Note that this will require either adding S | B | W | M to the list of constraints on the TypeVar or setting its bound to that.

[Daverball]

I.e. define a generic abstract base class Sliding and then define a concrete generic class for each of the four modes SlidingWindows SlidingSlices SlidingBounds and SlidingsMasks.

[Daverball]

You could even go one step further and define a __new__ on Sliding to return the correct subclass based on the mode parameter.

[randolf-scholz]

First of, thanks for the thoughtful response, it seems like this is a way to get at least pyright happy. mypy still seems to suffer from the same bug...

since the different modes don't actually have that much shared code.

Well, this is a simplified example. The concrete class has quite a bit more setup code in __init__, the difference between the modes is just in the return type of __iter__, in reality they share like 95% of code. I'll think about using separate classes, not sure how I feel about it right now.

[Daverball]

The shared code can still live in the base class, it would have an abstract __iter__ method in this case and each mode implements just that.

I meant shared more in the sense that there aren't complex conditionals where you would need to introduce multiple levels of inheritance to keep code reuse low.

[erictraut]

I've included a full response in the pyright issue. I'll paste the portion of my response that is relevant here and may shed some light on what's going on.

The overload matching logic in both pyright and mypy have special handling for Any. In particular, if Any appears in an argument passed to a call and this results in an ambiguous overload selection, both type checkers will evaluate the return type of the call as Any. However, in your case, this isn't applicable because the Any appears only in the self argument, and this argument is eliminated from the overload signatures during binding (i.e. during the evaluation of the sampler.__iter__ expression, before the call expression is evaluated). By the time the overload matching logic is invoked, there is no Any argument present. The type checker is simply choosing from among the following overloads:

    @overload
    def __iter__() -> Iterator[tuple[T, T]]: ...
    @overload
    def __iter__() -> Iterator[list[T]]: ...
    @overload
    def __iter__() -> Iterator[slice]: ...

This is why both mypy and pyright choose the first overload in this case.