Pyright: Invalid inference for zip(*a) (zip with asterisk unpacking)

[tamuhey]

Describe the bug

l: List[Tuple[int, str]] = [(1, "a"), (2, "b")]
b: Tuple[str]
a, b = zip(*l) # Error: Expression of type 'Tuple[int]' cannot be assigned to declared type 'Tuple[str]'
print(b)

('a', 'b')

The type of b is Tuple[str], but pyright infers Tuple[int]

Expected behavior

No error

[erictraut]

Thanks for the bug report.

There is a bug in Pyright here, but even with a fix, an error will still be generated here, so I'm not able to really solve your problem.

Here are the details, if you're interested.

When the type checker sees the expression zip(*l), it attempts to match the argument *l against the first parameter of zip. The zip function is defined in the builtins.pyi type stub as an overloaded method that takes one to six parameters, all of them iterables. Here's how it defines the one-parameter version:

    @overload
    def zip(__iter1: Iterable[_T1]) -> Iterator[Tuple[_T1]]: ...

In your example, the expression *l uses an unpack operator. But a type checker has no way of knowing how many elements will be unpacked because the type of l doesn't encode its length. That means the type checker can't determine which version of zip to use. It chooses the first one that potentially matches, which is the one-parameter version.

The type of the argument expression *l is Tuple[int, str], so the type checker matches Tuple[int, str] to the parameter type Iterable[_T1]. Here's where Pyright contains a bug. It doesn't correctly handle heterogeneous tuples when matching against a generic protocol class like Iterable. Because of this bug, it assigns the TypeVar _T1 the type int. The resulting return type for zip is therefore determined to be Iterator[Tuple[int]].

The correct behavior in this case is to create a union out of all element types in the heterogeneous Tuple, so the TypeVar _T1 should be assigned the type Union[int, str]. That makes sense since iterating over a tuple with int and str elements will result in both types, hence the union.

With this bug fixed, the expected return result for zip is Iterator[Tuple[Union[int, str]], and that's not going to be assignable to the expression a, b.

There are two problems here. Both of them are due to limitations in the expressiveness of Python's type extensions. The first problem is one that I already mentioned above with the unpack operator. A type checker has no way to determine statically how many parameters will be matched by an unpacked iterable, so it can't determine which overloaded version of zip it should use.

The second problem is that the generic type Iterable loses information about the ordering of elements within a Tuple, which is why we need to create a union out of the element types.

Possible workarounds:

1. Avoid using zip in this manner.
2. Add a "# type: ignore" comment to silence the type checker on that line.
3. Add an explicit type cast.

[erictraut]

The bug I mentioned above will be fixed in the next version of pyright.

[erictraut]

This is fixed in Pyright 1.1.22, which I just published.