Skip to content

Commit bed49ab

Browse files
authored
Fix unpack with overloaded __iter__ method (#14817)
The change is quite simple: instead of manually unpacking the return type of `find_item`, we use the regular machinery for it. Note, that this PR does not include any metaclasses related work. We still have problems with them, but I will defer them to the next PR. Closes #14811
1 parent 90b9a1a commit bed49ab

File tree

2 files changed

+73
-8
lines changed

2 files changed

+73
-8
lines changed

mypy/checker.py

Lines changed: 7 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -6395,14 +6395,13 @@ def iterable_item_type(self, instance: Instance) -> Type:
63956395
# in case there is no explicit base class.
63966396
return item_type
63976397
# Try also structural typing.
6398-
iter_type = get_proper_type(find_member("__iter__", instance, instance, is_operator=True))
6399-
if iter_type and isinstance(iter_type, CallableType):
6400-
ret_type = get_proper_type(iter_type.ret_type)
6401-
if isinstance(ret_type, Instance):
6402-
iterator = map_instance_to_supertype(
6403-
ret_type, self.lookup_typeinfo("typing.Iterator")
6404-
)
6405-
item_type = iterator.args[0]
6398+
ret_type, _ = self.expr_checker.check_method_call_by_name(
6399+
"__iter__", instance, [], [], instance
6400+
)
6401+
ret_type = get_proper_type(ret_type)
6402+
if isinstance(ret_type, Instance):
6403+
iterator = map_instance_to_supertype(ret_type, self.lookup_typeinfo("typing.Iterator"))
6404+
item_type = iterator.args[0]
64066405
return item_type
64076406

64086407
def function_type(self, func: FuncBase) -> FunctionLike:

test-data/unit/check-classes.test

Lines changed: 66 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -4524,6 +4524,72 @@ WithMeta().a # E: "WithMeta" has no attribute "a"
45244524
t: Type[WithMeta]
45254525
t.unknown # OK
45264526

4527+
[case testUnpackIterableClassWithOverloadedIter]
4528+
from typing import Generic, overload, Iterator, TypeVar, Union
4529+
4530+
AnyNum = TypeVar('AnyNum', int, float)
4531+
4532+
class Foo(Generic[AnyNum]):
4533+
@overload
4534+
def __iter__(self: Foo[int]) -> Iterator[float]: ...
4535+
@overload
4536+
def __iter__(self: Foo[float]) -> Iterator[int]: ...
4537+
def __iter__(self) -> Iterator[Union[float, int]]:
4538+
...
4539+
4540+
a, b, c = Foo[int]()
4541+
reveal_type(a) # N: Revealed type is "builtins.float"
4542+
reveal_type(b) # N: Revealed type is "builtins.float"
4543+
reveal_type(c) # N: Revealed type is "builtins.float"
4544+
4545+
x, y = Foo[float]()
4546+
reveal_type(x) # N: Revealed type is "builtins.int"
4547+
reveal_type(y) # N: Revealed type is "builtins.int"
4548+
[builtins fixtures/list.pyi]
4549+
4550+
[case testUnpackIterableClassWithOverloadedIter2]
4551+
from typing import Union, TypeVar, Generic, overload, Iterator
4552+
4553+
X = TypeVar('X')
4554+
4555+
class Foo(Generic[X]):
4556+
@overload
4557+
def __iter__(self: Foo[str]) -> Iterator[int]: ... # type: ignore
4558+
@overload
4559+
def __iter__(self: Foo[X]) -> Iterator[str]: ...
4560+
def __iter__(self) -> Iterator[Union[int, str]]:
4561+
...
4562+
4563+
a, b, c = Foo[str]()
4564+
reveal_type(a) # N: Revealed type is "builtins.int"
4565+
reveal_type(b) # N: Revealed type is "builtins.int"
4566+
reveal_type(c) # N: Revealed type is "builtins.int"
4567+
4568+
x, y = Foo[float]()
4569+
reveal_type(x) # N: Revealed type is "builtins.str"
4570+
reveal_type(y) # N: Revealed type is "builtins.str"
4571+
[builtins fixtures/list.pyi]
4572+
4573+
[case testUnpackIterableRegular]
4574+
from typing import TypeVar, Generic, Iterator
4575+
4576+
X = TypeVar('X')
4577+
4578+
class Foo(Generic[X]):
4579+
def __iter__(self) -> Iterator[X]:
4580+
...
4581+
4582+
a, b = Foo[int]()
4583+
reveal_type(a) # N: Revealed type is "builtins.int"
4584+
reveal_type(b) # N: Revealed type is "builtins.int"
4585+
[builtins fixtures/list.pyi]
4586+
4587+
[case testUnpackNotIterableClass]
4588+
class Foo: ...
4589+
4590+
a, b, c = Foo() # E: "Foo" object is not iterable
4591+
[builtins fixtures/list.pyi]
4592+
45274593
[case testMetaclassIterable]
45284594
from typing import Iterable, Iterator
45294595

0 commit comments

Comments
 (0)