FailureGod
FailureGod

Reputation: 382

Make mypy recognize asserting type of list

I have code like this

from typing import Union, List

class Player:
    number: str

def func(things: List[Union[Player, str]]):
    if isinstance(things[0], Player):
        print(" ".join(p.number for p in things))
    else:
        print(" ".join(things))

Mypy highlights the p.number and the things in the else block gives me this error:

[mypy error] [E] Item "str" of "Union[Player, str]" has no attribute "number"

I have also tried

def func2(things: List[Union[Player, str]]):
    if all(isinstance(thing, Player) for thing in things):
        print(" ".join(p.number for p in things))
    elif all(isinstance(thing, str) for thing in things):
        print(" ".join(things))

But I get the same error. How do I get mypy to recognize that I'm asserting that every element of a list is a specific type?

Upvotes: 3

Views: 1276

Answers (1)

Samwise
Samwise

Reputation: 71424

You don't want to have a "list of (players or strings)", you want "(a list of players) or (a list of strings)":

def func(things: Union[List[Player], List[str]]):
    if isinstance(things[0], Player):
        things = cast(List[Player], things)
        print(" ".join(p.number for p in things))
    else:
        things = cast(List[str], things)
        print(" ".join(things))

I don't think there's a more graceful way to get the typechecking to work than the cast, though; mypy doesn't seem to be quite clever enough to infer the type of things based on your things[0] check alone, and you can't do isinstance(things, List[Player]).

Upvotes: 4

Related Questions