From 3cc25f23e4c4d188e1dedb4a7c9c1506b4c8bb69 Mon Sep 17 00:00:00 2001 From: sobolevn Date: Thu, 2 Mar 2023 10:05:43 +0300 Subject: [PATCH 1/2] Fix unpack with overloaded `__iter__` method --- mypy/checker.py | 14 +++---- test-data/unit/check-classes.test | 66 +++++++++++++++++++++++++++++++ 2 files changed, 72 insertions(+), 8 deletions(-) diff --git a/mypy/checker.py b/mypy/checker.py index bf9905e7d8ee..4b0de6d7341c 100644 --- a/mypy/checker.py +++ b/mypy/checker.py @@ -6395,14 +6395,12 @@ def iterable_item_type(self, instance: Instance) -> Type: # in case there is no explicit base class. return item_type # Try also structural typing. - iter_type = get_proper_type(find_member("__iter__", instance, instance, is_operator=True)) - if iter_type and isinstance(iter_type, CallableType): - ret_type = get_proper_type(iter_type.ret_type) - if isinstance(ret_type, Instance): - iterator = map_instance_to_supertype( - ret_type, self.lookup_typeinfo("typing.Iterator") - ) - item_type = iterator.args[0] + ret_type, _ = self.expr_checker.check_method_call_by_name("__iter__", instance, [], [], instance) + if isinstance(ret_type, Instance): + iterator = map_instance_to_supertype( + ret_type, self.lookup_typeinfo("typing.Iterator") + ) + item_type = iterator.args[0] return item_type def function_type(self, func: FuncBase) -> FunctionLike: diff --git a/test-data/unit/check-classes.test b/test-data/unit/check-classes.test index d5fb830487e8..da95674ed08f 100644 --- a/test-data/unit/check-classes.test +++ b/test-data/unit/check-classes.test @@ -4524,6 +4524,72 @@ WithMeta().a # E: "WithMeta" has no attribute "a" t: Type[WithMeta] t.unknown # OK +[case testUnpackIterableClassWithOverloadedIter] +from typing import Generic, overload, Iterator, TypeVar, Union + +AnyNum = TypeVar('AnyNum', int, float) + +class Foo(Generic[AnyNum]): + @overload + def __iter__(self: Foo[int]) -> Iterator[float]: ... + @overload + def __iter__(self: Foo[float]) -> Iterator[int]: ... + def __iter__(self) -> Iterator[Union[float, int]]: + ... + +a, b, c = Foo[int]() +reveal_type(a) # N: Revealed type is "builtins.float" +reveal_type(b) # N: Revealed type is "builtins.float" +reveal_type(c) # N: Revealed type is "builtins.float" + +x, y = Foo[float]() +reveal_type(x) # N: Revealed type is "builtins.int" +reveal_type(y) # N: Revealed type is "builtins.int" +[builtins fixtures/list.pyi] + +[case testUnpackIterableClassWithOverloadedIter2] +from typing import Union, TypeVar, Generic, overload, Iterator + +X = TypeVar('X') + +class Foo(Generic[X]): + @overload + def __iter__(self: Foo[str]) -> Iterator[int]: ... # type: ignore + @overload + def __iter__(self: Foo[X]) -> Iterator[str]: ... + def __iter__(self) -> Iterator[Union[int, str]]: + ... + +a, b, c = Foo[str]() +reveal_type(a) # N: Revealed type is "builtins.int" +reveal_type(b) # N: Revealed type is "builtins.int" +reveal_type(c) # N: Revealed type is "builtins.int" + +x, y = Foo[float]() +reveal_type(x) # N: Revealed type is "builtins.str" +reveal_type(y) # N: Revealed type is "builtins.str" +[builtins fixtures/list.pyi] + +[case testUnpackIterableRegular] +from typing import TypeVar, Generic, Iterator + +X = TypeVar('X') + +class Foo(Generic[X]): + def __iter__(self) -> Iterator[X]: + ... + +a, b = Foo[int]() +reveal_type(a) # N: Revealed type is "builtins.int" +reveal_type(b) # N: Revealed type is "builtins.int" +[builtins fixtures/list.pyi] + +[case testUnpackNotIterableClass] +class Foo: ... + +a, b, c = Foo() # E: "Foo" object is not iterable +[builtins fixtures/list.pyi] + [case testMetaclassIterable] from typing import Iterable, Iterator From 1500162708d9c54e31b73568d26734a57bfb2182 Mon Sep 17 00:00:00 2001 From: sobolevn Date: Thu, 2 Mar 2023 10:20:41 +0300 Subject: [PATCH 2/2] Fix CI --- mypy/checker.py | 9 +++++---- 1 file changed, 5 insertions(+), 4 deletions(-) diff --git a/mypy/checker.py b/mypy/checker.py index 4b0de6d7341c..a913d2a8eb16 100644 --- a/mypy/checker.py +++ b/mypy/checker.py @@ -6395,11 +6395,12 @@ def iterable_item_type(self, instance: Instance) -> Type: # in case there is no explicit base class. return item_type # Try also structural typing. - ret_type, _ = self.expr_checker.check_method_call_by_name("__iter__", instance, [], [], instance) + ret_type, _ = self.expr_checker.check_method_call_by_name( + "__iter__", instance, [], [], instance + ) + ret_type = get_proper_type(ret_type) if isinstance(ret_type, Instance): - iterator = map_instance_to_supertype( - ret_type, self.lookup_typeinfo("typing.Iterator") - ) + iterator = map_instance_to_supertype(ret_type, self.lookup_typeinfo("typing.Iterator")) item_type = iterator.args[0] return item_type