How to narrow a list of unions?

user_3067276 0
Viewed 118

I have a variable of type list[A | B] that could hold a mixed list (like [A(), B(), A()]).

If I later reach some corner case and I want to make sure all elements are actually of type A, I can assert isinstance inside a for-loop:

def f(mylist: list[A | B]):
    ...
    ...
    for el in mylist:
        assert isinstance(el, A)
    # Now, I'm sure that mylist is actually `list[A]`
    # How do I tell that to the type checker?

After the for-loop, if I reveal_type(mylist) it still says list[A|B]. I also tried assert all(isinstance(el, A) for el in mylist) instead of the explicit loop, but mypy still isn't able to narrow it. Is this possible? Or do I have to use cast here?

narrowing union-types python-typing mypy python
1 Answers

Or do I have to use cast here?

Using Python 3.10 the only alternative to using cast that can narrow a list would be using a TypeGuard, this code:

from typing import TypeGuard


def is_str_list(val: list[int | str]) -> TypeGuard[list[str]]:
    """Determines whether all objects in the list are strings"""
    return all(isinstance(x, str) for x in val)


def is_int_list(val: list[int | str]) -> TypeGuard[list[int]]:
    """Determines whether all objects in the list are ints"""
    return all(isinstance(x, int) for x in val)


def f(mylist: list[int | str]):

    if is_str_list(mylist):
        reveal_type(mylist)

    elif is_int_list(mylist):
        reveal_type(mylist)

    reveal_type(mylist)

narrows the type of the list without using cast:

narrow_collection.py:17: note: Revealed type is "builtins.list[builtins.str]"
narrow_collection.py:21: note: Revealed type is "builtins.list[builtins.int]"
narrow_collection.py:24: note: Revealed type is "builtins.list[Union[builtins.int, builtins.str]]"
user_107940310
Related