python
type annotated arguments in zip function lead to wrong return types when using the star operator
This issue is a followup to the following issue: https://github.com/microsoft/pylance-release/issues/3598
Here's the code that is wrongly annotated:
test_list: list[tuple[int, str]] = [(1, 'testa'), (2, 'testb')]
a: tuple[int]
b: tuple[str]
a, b = zip(*test_list)
Pylance warns about an assignment of typle[int | str] to the tuples a and b. The tuple annotation is being converted to an Iterator annotation, which does not support a sequence of types (see Eric's reply to the previous issue for details). Therefore, the resulting tuples after the unpacking operation are not as precisely typed as they could be.
And here is the typeshed code that leads to the 'wrong' type annotations:
https://github.com/python/typeshed/blob/7b3fff714a48f0413141ba10344f9473e3ec4a18/stdlib/builtins.pyi#L1673-L1715
I'm not sure if it is currently possible to fix this. I had the idea of adding a type hint to the *args keyword, so something like this:
@overload
def __new__(cls, __iter1: _T2:=Sequence[_T1], *args: _T2, strict: bool = ...) -> zip[tuple[_T1]]:
I'm unsure if this is valid python syntax and if it would result in the desired behavior. I'm also unsure how python would decide between the proposed constructor and the one that is currently being used during typechecking my example: https://github.com/python/typeshed/blob/7b3fff714a48f0413141ba10344f9473e3ec4a18/stdlib/builtins.pyi#L1675-L1676
This code works for me:
from typing_extensions import assert_type
test_list: list[tuple[int, str]] = [(1, 'testa'), (2, 'testb')]
a: tuple[int, ...] # note
b: tuple[str, ...]
a, b = zip(*test_list)
assert_type(a, tuple[int, ...])
assert_type(b, tuple[str, ...])
However, it does not work without explicit a and b annotations:
Running mypy --platform darwin --python-version 3.10 on the standard library test cases... failure
test_cases/stdlib/builtins/check_zip.py:9: error: Expression is of type "Tuple[Any, ...]", not "Tuple[int, ...]" [assert-type]
test_cases/stdlib/builtins/check_zip.py:10: error: Expression is of type "Tuple[Any, ...]", not "Tuple[str, ...]" [assert-type]
A reduced version of the problem is:
test_list: list[tuple[int, str]] = []
reveal_type(zip(*test_list))
Mypy infers this as zip[tuple[Any, ...]], pyright as zip[tuple[int | str]]. Mypy's answer works better here, but pyright's overload heuristic has other advantages and probably this cannot be fixed in pyright.
Ideally we'd write the stub in such a way that it works well for all major type checkers, but I'm not sure how to do that in this case without regressing support for other important use cases.
Is zip not missing an overload here?
m: list[Iterable[int]] = [[42,2],[4,5]]
a = zip(*m) # reveal_type zip[tuple[Any, ...]]
I expect a: zip[tuple[int, ...]] instead of
I think we are missing one overload:
@overload
def __new__[_T1](
cls,
*iterables: Iterable[_T1],
strict: bool = ...,
) -> zip[tuple[_T1, ...]]: ...
