Recognise list comprehension type

Question

Take this following code:

from typing import Union

value: Union[str, bytes]
stripped = [line.strip() for line in value.splitlines()]

reveal_type(stripped)

To anyone reading, it's clear that stripped should be of type

Union[List[str], List[bytes]]

However, mypy recognised it as

List[Union[str, bytes]]

How can I get mypy to recognise it as Union[List[str], List[bytes]]?

Answer 1

You could try this:

from typing import Union

value: Union[str, bytes]
stripped_str =[]

if isinstance(value, str):
    stripped_str = [line.strip() for line in value.splitlines()]
else:
    stripped_bytes = [line.strip() for line in value.splitlines()]

stripped = stripped_str or stripped_bytes

reveal_type(stripped)

So that Mypy finds the type of stripped is Union[list[str], list[bytes]].