Issues with unions of sequences

[gsakkis]

Consider the following example:

import random
from typing import TYPE_CHECKING, Sequence, Union, overload


@overload
def double(n: int) -> int:
    ...

@overload
def double(n: str) -> str:
    ...

def double(n):
    return 2 * n


x: Sequence[int] = [1, 2, 3]
y: Sequence[str] = ["a", "b", "c"]
z = x if random.random() > 0.5 else y
if TYPE_CHECKING:
    reveal_type(z)
    reveal_type(z[0])
    reveal_type(double(z[0]))
    reveal_type(iter(z))

z2 = map(double, z)
if TYPE_CHECKING:
    reveal_type(z2)
    reveal_type(next(z2))

❌ 1a. Without an explicit annotation for z, its inferred type is Sequence[object] instead of Union[Sequence[int], Sequence[str]]:

reveal_type(z)
note: Revealed type is "typing.Sequence[builtins.object]"

reveal_type(z[0])
note: Revealed type is "builtins.object"

reveal_type(double(z[0]))
error: No overload variant of "double" matches argument type "object"  [call-overload]
note: Possible overload variants:
note:     def double(n: int) -> int
note:     def double(n: str) -> str

reveal_type(iter(z))
note: Revealed type is "Any"

❌ 1b. A map over an overloaded function seems to take into account only the first overload:

z2 = map(double, z)
error: Argument 1 to "map" has incompatible type overloaded function; expected "Callable[[object], int]"  [arg-type]

reveal_type(z2)
note: Revealed type is "builtins.map[builtins.int]"

reveal_type(next(z2))
note: Revealed type is "builtins.int"

✔️ 2a. Annotating z explicitly fixes some issues by preserving the Union:

z: Union[Sequence[int], Sequence[str]] = x if random.random() > 0.5 else y
z2 = map(double, z)

reveal_type(z)
note: Revealed type is "Union[typing.Sequence[builtins.int], typing.Sequence[builtins.str]]"

reveal_type(z[0])
note: Revealed type is "Union[builtins.int, builtins.str]"

reveal_type(double(z[0]))
note: Revealed type is "Union[builtins.int, builtins.str]"

reveal_type(z2)
note: Revealed type is "Union[builtins.map[builtins.int], builtins.map[builtins.str]]"

❌ 2b. But not all of them; iter and next fall back to object:

reveal_type(iter(z))
note: Revealed type is "typing.Iterator[builtins.object]"

reveal_type(next(z2))
note: Revealed type is "builtins.object"

[AlexWaygood]

This is a classic of the "join versus union" genre of mypy bugs. It's similar to #8586, #12327, #11440, #10135, etc. #12056 is the meta-issue for this genre of bugs.

This is a valid bug report, but I think we have enough bug reports in this genre open already, and we know what the problem is here, so I'm closing this.

[gsakkis]

Thanks for looking into it, indeed (1a) and (1b) above fall under the join versus union issue. For (2b) though z is explicitly annotated as a union of sequences and indexing preserves the union, only iter and next don't. Could this be an independent issue?

[AlexWaygood]

Thanks for looking into it, indeed (1a) and (1b) above fall under the join versus union issue. For (2b) though z is explicitly annotated as a union of sequences and indexing preserves the union, only iter and next don't. Could this be an independent issue?

I think 2b) is still a union-vs-join issue. The issue is that an object of type Union[Sequence[int], Sequence[str]] is being fed into the iter() function. Mypy looks at the stub for iter() in typeshed, and realises it doesn't know how to infer the return type, since there's no overload for a union input in the stub:

@overload
def iter(__iterable: SupportsIter[_SupportsNextT]) -> _SupportsNextT: ...
@overload
def iter(__iterable: _GetItemIterable[_T]) -> Iterator[_T]: ...
@overload
def iter(__function: Callable[[], _T | None], __sentinel: None) -> Iterator[_T]: ...
@overload
def iter(__function: Callable[[], _T], __sentinel: object) -> Iterator[_T]: ...

In order to infer the return type, therefore, I think mypy then performs a join operation to transform Union[Sequence[int], Sequence[str]] into Sequence[object]. Once it does that, it finds the first overload matches, and infers the return type as Iterator[object].

[AlexWaygood]

Observe the difference between these two snippets:

from typing import Sequence, Union
import random

x: Sequence[int] = [1, 2, 3]
y: Sequence[str] = ["a", "b", "c"]

z: Union[Sequence[int], Sequence[str]] = x if random.random() > 0.5 else y
a: Sequence[Union[int, str]] = x if random.random() > 0.5 else y  # type: ignore

reveal_type(iter(z))  # object
reveal_type(iter(a))  # Union[int, str]

https://mypy-play.net/?mypy=latest&python=3.10&gist=41cd61709ac31f3c53096035db8eee3d

[gsakkis]

Ok I see. I subscribed to the meta-issue, when there's a resolution I'll try this again and reopen the issue if needed.