mypy fails to narrow down generic types in `TypeGuard`-based conditionals

Question

Trying to create a type-pure container in Python, I've stumbled upon type-narrowing failure of generic types by mypy when creating a signature of the __setitem__ method.

Before we dive in, I'm using:

• mypy 0.991
• python 3.10.6.

---

First, let's consider a basic example.

Let us define two types, A and B, and a type variable AB that can be either A or B (but not a Union[A, B]):

from __future__ import annotations

import typing as t
from collections import abc
from itertools import tee
from typing_extensions import reveal_type


class A:
    ...
    

class B:
    ...

AB = t.TypeVar('AB', A, B)  # Strictly A or B

Now, suppose we want to define a function or method with the following signature:

def fn(arg: AB | abc.Iterable[AB]):

The function's body will depend on whether AB or abc.Iterable[AB] are provided. To check that, let's introduce the following utilities:

T = t.TypeVar('T')


def is_iterable_of(
        s: abc.Iterable[t.Any], _type: t.Type[T]
) -> t.TypeGuard[abc.Iterable[T]]:
    return all(isinstance(x, _type) for x in s)


def is_type(x: t.Any, _type: t.Type[T]) -> t.TypeGuard[T]:
    return isinstance(x, _type)


def is_a_or_b(x: t.Any) -> t.TypeGuard[A | B]:
    return isinstance(x, A) or isinstance(x, B)


def is_iterable_of_ab(s: abc.Iterable[t.Any]) -> t.TypeGuard[abc.Iterable[AB]]:
    ss = tee(s)
    
    return any(
            is_iterable_of(_s, _t) for _s, _t in zip(ss, [A, B])
        )

I think only is_iterable_of_ab might be a little confusing: it checks whether the provided iterable contains all A types or all B types. To view what TypeGuard types are doing, please refer to the mypy docs.

We apply our type-narrowing functions in fn as follows:

def fn(arg: A | B | abc.Iterable[AB]):
    if is_a_or_b(arg):
        # Process single instance
        reveal_type(arg)  # A | B
    else:
        # Process multiple instances
        reveal_type(arg)  # (A | B | Iterable[A]) | (A | B | Iterable[B])
        assert is_iterable_of_ab(arg)

Mypy produces the following output (you can check the full example's code in mypy playground):

main.py:46: note: Revealed type is "Union[__main__.A, __main__.B]"
main.py:49: note: Revealed type is "Union[__main__.A, __main__.B, typing.Iterable[__main__.A]]"
main.py:49: note: Revealed type is "Union[__main__.A, __main__.B, typing.Iterable[__main__.B]]"
main.py:50: error: Argument 1 to "is_iterable_of_ab" has incompatible type "Union[A, B, Iterable[A]]"; expected "Iterable[Any]"  [arg-type]
main.py:50: error: Argument 1 to "is_iterable_of_ab" has incompatible type "Union[A, B, Iterable[B]]"; expected "Iterable[Any]"  [arg-type]

What is confusing here is that the call to is_a_or_b, while narrows down the type to A | B correctly, doesn't lead to mypy recognizing that if the conditional is False, the provided arg must have the abc.Iterable[AB] type. As a result, any functionality inside the else block that depends on arg being Iterable[AB] will raise an error since mypy thinks arg is can still be A | B.

Note that this also fails:

def fn2(arg: A | B | abc.Iterable[AB]):
    if is_a_or_b(arg):
        # Process single instance
        reveal_type(arg)  # A | B
    elif is_iterable_of_ab(arg):  # error: incompatible type
        # Process multiple instances
        reveal_type(arg)  # wrong type: Iterable[A]
    else:
        raise TypeError()

So, how could I correctly use type-narrowing here?

---

For a more realistic example indicated at the beginning, consider making a type pure list class:

class TypePureList(t.Generic[AB]):
    
    def __init__(self, items: abc.Iterable[AB]):
        if not isinstance(items, list):
            items = list(items)
        self._items: list[AB] = items
    
    @t.overload
    def __setitem__(self, index: t.SupportsIndex, value: AB):
        ...
    
    @t.overload
    def __setitem__(self, index: slice, value: abc.Iterable[AB]):
        ...
    
    def __setitem__(self, index: t.SupportsIndex | slice, value: AB | abc.Iterable[AB]):
        
        if len(self._items) == 0:
            raise ValueError
        
        if isinstance(index, t.SupportsIndex):
            # Check that the value type is correct by itself
            assert is_a_or_b(value)
            # Check that the value type matches the existing type
            assert is_iterable_of_ab([self._items[0], value])
            reveal_type(value)  # A | B
        else:
            # Index must be slice
            # Value must be Iterable
            reveal_type(index)  # slice
            reveal_type(value)  # (A | Iterable[A]) | (B | Iterable[B])
            # Check that the iterable contains correct types
            value, value1, value2 = tee(value, 3)  # error: incompatible type
            assert is_iterable_of_ab(value1)
            # Check if the iterable types match the existing type
            assert is_iterable_of_ab([self._items[0], next(value2)])

        # error: invalid syntax -- index is Union[SupportsIndex, slice]
        # error: incompatible types in assignment
        self._items.__setitem__(index, value)

Firstly, mypy fails to use the overloading information, but that seems to be an ongoing issue (see #7858). Secondly, same as above, mypy fails to narrow down the types following the assertion that if index has SupportsIndex type, then value must be Iterable[AB] due to assert is_a_or_b(value).

---