From b212702ddb0e77a3b48759a4f5fdc1498356f035 Mon Sep 17 00:00:00 2001 From: AlexWaygood Date: Fri, 3 Mar 2023 12:23:50 +0000 Subject: [PATCH 1/3] Allow iterable class objects to be unpacked (including enums) --- mypy/checker.py | 21 +++--- test-data/unit/check-inference.test | 102 ++++++++++++++++++++++++++++ test-data/unit/pythoneval.test | 17 +++++ 3 files changed, 131 insertions(+), 9 deletions(-) diff --git a/mypy/checker.py b/mypy/checker.py index 5f72e6f7f399d..ed95e087c406b 100644 --- a/mypy/checker.py +++ b/mypy/checker.py @@ -3632,7 +3632,9 @@ def check_multi_assignment_from_iterable( infer_lvalue_type: bool = True, ) -> None: rvalue_type = get_proper_type(rvalue_type) - if self.type_is_iterable(rvalue_type) and isinstance(rvalue_type, Instance): + if self.type_is_iterable(rvalue_type) and isinstance( + rvalue_type, (Instance, CallableType, TypeType) + ): item_type = self.iterable_item_type(rvalue_type) for lv in lvalues: if isinstance(lv, StarExpr): @@ -6387,15 +6389,16 @@ def note( return self.msg.note(msg, context, offset=offset, code=code) - def iterable_item_type(self, instance: Instance) -> Type: - iterable = map_instance_to_supertype(instance, self.lookup_typeinfo("typing.Iterable")) - item_type = iterable.args[0] - if not isinstance(get_proper_type(item_type), AnyType): - # This relies on 'map_instance_to_supertype' returning 'Iterable[Any]' - # in case there is no explicit base class. - return item_type + def iterable_item_type(self, it: Instance | CallableType | TypeType) -> Type: + if isinstance(it, Instance): + iterable = map_instance_to_supertype(it, self.lookup_typeinfo("typing.Iterable")) + item_type = iterable.args[0] + if not isinstance(get_proper_type(item_type), AnyType): + # This relies on 'map_instance_to_supertype' returning 'Iterable[Any]' + # in case there is no explicit base class. + return item_type # Try also structural typing. - return self.analyze_iterable_item_type_without_expression(instance, instance)[1] + return self.analyze_iterable_item_type_without_expression(it, it)[1] def function_type(self, func: FuncBase) -> FunctionLike: return function_type(func, self.named_type("builtins.function")) diff --git a/test-data/unit/check-inference.test b/test-data/unit/check-inference.test index cfb553820d9ef..6acb5804577ce 100644 --- a/test-data/unit/check-inference.test +++ b/test-data/unit/check-inference.test @@ -270,6 +270,108 @@ def f() -> None: class A: pass [out] +[case testClassObjectsNotUnpackableWithoutIterableMetaclass] +from typing import Type + +class Foo: ... +A: Type[Foo] = Foo +a, b = Foo # E: "Type[Foo]" object is not iterable +c, d = A # E: "Type[Foo]" object is not iterable + +class Meta(type): ... +class Bar(metaclass=Meta): ... +B: Type[Bar] = Bar +e, f = Bar # E: "Type[Bar]" object is not iterable +g, h = B # E: "Type[Bar]" object is not iterable + +reveal_type(a) # E: Cannot determine type of "a" # N: Revealed type is "Any" +reveal_type(b) # E: Cannot determine type of "b" # N: Revealed type is "Any" +reveal_type(c) # E: Cannot determine type of "c" # N: Revealed type is "Any" +reveal_type(d) # E: Cannot determine type of "d" # N: Revealed type is "Any" +reveal_type(e) # E: Cannot determine type of "e" # N: Revealed type is "Any" +reveal_type(f) # E: Cannot determine type of "f" # N: Revealed type is "Any" +reveal_type(g) # E: Cannot determine type of "g" # N: Revealed type is "Any" +reveal_type(h) # E: Cannot determine type of "h" # N: Revealed type is "Any" +[out] + +[case testInferringLvarTypesUnpackedFromIterableClassObject] +from typing import Iterator, Type, TypeVar, Union +class Meta(type): + def __iter__(cls) -> Iterator[int]: + yield from [1, 2, 3] + +class Meta2(type): + def __iter__(cls) -> Iterator[str]: + yield from ["foo", "bar", "baz"] + +class Meta3(type): ... + +class Foo(metaclass=Meta): ... +class Bar(metaclass=Meta2): ... +class Baz(metaclass=Meta3): ... +class Spam: ... + +A: Type[Foo] = Foo +B: Type[Union[Foo, Bar]] = Foo +C: Union[Type[Foo], Type[Bar]] = Foo +D: Type[Union[Foo, Baz]] = Foo +E: Type[Union[Foo, Spam]] = Foo + +a, b, c = Foo +d, e, f = A +g, h, i = B +j, k, l = C +m, n, o = D # E: "Type[Baz]" object is not iterable +p, q, r = E # E: "Type[Spam]" object is not iterable + +for var in [a, b, c, d, e, f]: + reveal_type(var) # N: Revealed type is "builtins.int" + +for var2 in [g, h, i, j, k, l]: + reveal_type(var2) # N: Revealed type is "Union[builtins.int, builtins.str]" + +for var3 in [m, n, o, p, q, r]: + reveal_type(var3) # N: Revealed type is "Union[builtins.int, Any]" + +T = TypeVar("T", bound=Type[Foo]) + +def check(x: T) -> T: + a, b, c = x + for var in [a, b, c]: + reveal_type(var) # N: Revealed type is "builtins.int" + return x + +T2 = TypeVar("T2", bound=Type[Union[Foo, Bar]]) + +def check2(x: T2) -> T2: + a, b, c = x + for var in [a, b, c]: + reveal_type(var) # N: Revealed type is "Union[builtins.int, builtins.str]" + return x + +T3 = TypeVar("T3", bound=Union[Type[Foo], Type[Bar]]) + +def check3(x: T3) -> T3: + a, b, c = x + for var in [a, b, c]: + reveal_type(var) # N: Revealed type is "Union[builtins.int, builtins.str]" + return x +[out] + +[case testInferringLvarTypesUnpackedFromIterableClassObjectWithGenericIter] +from typing import Iterator, Type, TypeVar + +T = TypeVar("T") +class Meta(type): + def __iter__(self: Type[T]) -> Iterator[T]: ... +class Foo(metaclass=Meta): ... + +A, B, C = Foo +reveal_type(A) # N: Revealed type is "__main__.Foo" +reveal_type(B) # N: Revealed type is "__main__.Foo" +reveal_type(C) # N: Revealed type is "__main__.Foo" +[out] + [case testInferringLvarTypesInMultiDefWithInvalidTuple] from typing import Tuple t = None # type: Tuple[object, object, object] diff --git a/test-data/unit/pythoneval.test b/test-data/unit/pythoneval.test index a3413e0711849..d89009c2c56f5 100644 --- a/test-data/unit/pythoneval.test +++ b/test-data/unit/pythoneval.test @@ -1878,6 +1878,23 @@ _testEnumIterMetaInference.py:8: note: Revealed type is "typing.Iterator[_E`-1]" _testEnumIterMetaInference.py:9: note: Revealed type is "_E`-1" _testEnumIterMetaInference.py:13: note: Revealed type is "socket.SocketKind" +[case testEnumUnpackedViaMetaclass] +from enum import Enum + +class FooEnum(Enum): + A = 1 + B = 2 + C = 3 + +a, b, c = FooEnum +reveal_type(a) +reveal_type(b) +reveal_type(c) +[out] +_testEnumUnpackedViaMetaclass.py:9: note: Revealed type is "_testEnumUnpackedViaMetaclass.FooEnum" +_testEnumUnpackedViaMetaclass.py:10: note: Revealed type is "_testEnumUnpackedViaMetaclass.FooEnum" +_testEnumUnpackedViaMetaclass.py:11: note: Revealed type is "_testEnumUnpackedViaMetaclass.FooEnum" + [case testNativeIntTypes] # Spot check various native int operations with full stubs. from mypy_extensions import i64, i32 From 0a1ecb291ffadc89408d5e61ab95384f020a2150 Mon Sep 17 00:00:00 2001 From: AlexWaygood Date: Fri, 3 Mar 2023 14:00:27 +0000 Subject: [PATCH 2/3] Handle overloaded `__init__` methods --- mypy/checker.py | 2 +- test-data/unit/check-inference.test | 16 ++++++++++++++-- 2 files changed, 15 insertions(+), 3 deletions(-) diff --git a/mypy/checker.py b/mypy/checker.py index ed95e087c406b..790871761124a 100644 --- a/mypy/checker.py +++ b/mypy/checker.py @@ -3633,7 +3633,7 @@ def check_multi_assignment_from_iterable( ) -> None: rvalue_type = get_proper_type(rvalue_type) if self.type_is_iterable(rvalue_type) and isinstance( - rvalue_type, (Instance, CallableType, TypeType) + rvalue_type, (Instance, CallableType, TypeType, Overloaded) ): item_type = self.iterable_item_type(rvalue_type) for lv in lvalues: diff --git a/test-data/unit/check-inference.test b/test-data/unit/check-inference.test index 6acb5804577ce..2dc19d319a0db 100644 --- a/test-data/unit/check-inference.test +++ b/test-data/unit/check-inference.test @@ -295,7 +295,7 @@ reveal_type(h) # E: Cannot determine type of "h" # N: Revealed type is "Any" [out] [case testInferringLvarTypesUnpackedFromIterableClassObject] -from typing import Iterator, Type, TypeVar, Union +from typing import Iterator, Type, TypeVar, Union, overload class Meta(type): def __iter__(cls) -> Iterator[int]: yield from [1, 2, 3] @@ -311,11 +311,20 @@ class Bar(metaclass=Meta2): ... class Baz(metaclass=Meta3): ... class Spam: ... +class Eggs(metaclass=Meta): + @overload + def __init__(self, x: int) -> None: ... + @overload + def __init__(self, x: int, y: int, z: int) -> None: ... + def __init__(self, x: int, y: int = ..., z: int = ...) -> None: ... + A: Type[Foo] = Foo B: Type[Union[Foo, Bar]] = Foo C: Union[Type[Foo], Type[Bar]] = Foo D: Type[Union[Foo, Baz]] = Foo E: Type[Union[Foo, Spam]] = Foo +F: Type[Eggs] = Eggs +G: Type[Union[Foo, Eggs]] = Foo a, b, c = Foo d, e, f = A @@ -323,8 +332,11 @@ g, h, i = B j, k, l = C m, n, o = D # E: "Type[Baz]" object is not iterable p, q, r = E # E: "Type[Spam]" object is not iterable +s, t, u = Eggs +v, w, x = F +y, z, aa = G -for var in [a, b, c, d, e, f]: +for var in [a, b, c, d, e, f, s, t, u, v, w, x, y, z, aa]: reveal_type(var) # N: Revealed type is "builtins.int" for var2 in [g, h, i, j, k, l]: From 2ee19697b3740753c42af0199c9e7554a8c9c3b3 Mon Sep 17 00:00:00 2001 From: AlexWaygood Date: Fri, 3 Mar 2023 14:03:07 +0000 Subject: [PATCH 3/3] Fix type hint --- mypy/checker.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/mypy/checker.py b/mypy/checker.py index 790871761124a..66facb3797f9e 100644 --- a/mypy/checker.py +++ b/mypy/checker.py @@ -6389,7 +6389,7 @@ def note( return self.msg.note(msg, context, offset=offset, code=code) - def iterable_item_type(self, it: Instance | CallableType | TypeType) -> Type: + def iterable_item_type(self, it: Instance | CallableType | TypeType | Overloaded) -> Type: if isinstance(it, Instance): iterable = map_instance_to_supertype(it, self.lookup_typeinfo("typing.Iterable")) item_type = iterable.args[0]