Skip to content

Commit 8af3af3

Browse files
authored
Support protocol inference for Type[T] via metaclass (#14554)
Fixes #12553 This looks quite niche, but also it was mentioned recently couple times for a real-life use case: enum classes, and implementation looks simple.
1 parent c4ecd2b commit 8af3af3

File tree

3 files changed

+51
-0
lines changed

3 files changed

+51
-0
lines changed

mypy/constraints.py

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -619,6 +619,17 @@ def visit_instance(self, template: Instance) -> list[Constraint]:
619619
actual.item, template, subtype, template, class_obj=True
620620
)
621621
)
622+
if self.direction == SUPERTYPE_OF:
623+
# Infer constraints for Type[T] via metaclass of T when it makes sense.
624+
a_item = actual.item
625+
if isinstance(a_item, TypeVarType):
626+
a_item = get_proper_type(a_item.upper_bound)
627+
if isinstance(a_item, Instance) and a_item.type.metaclass_type:
628+
res.extend(
629+
self.infer_constraints_from_protocol_members(
630+
a_item.type.metaclass_type, template, actual, template
631+
)
632+
)
622633

623634
if isinstance(actual, Overloaded) and actual.fallback is not None:
624635
actual = actual.fallback

test-data/unit/check-protocols.test

Lines changed: 21 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -3977,3 +3977,24 @@ class C:
39773977
DEFAULT: ClassVar[C]
39783978

39793979
x: P = C()
3980+
3981+
[case testInferenceViaTypeTypeMetaclass]
3982+
from typing import Iterator, Iterable, TypeVar, Type
3983+
3984+
M = TypeVar("M")
3985+
3986+
class Meta(type):
3987+
def __iter__(self: Type[M]) -> Iterator[M]: ...
3988+
class Foo(metaclass=Meta): ...
3989+
3990+
T = TypeVar("T")
3991+
def test(x: Iterable[T]) -> T: ...
3992+
3993+
reveal_type(test(Foo)) # N: Revealed type is "__main__.Foo"
3994+
t_foo: Type[Foo]
3995+
reveal_type(test(t_foo)) # N: Revealed type is "__main__.Foo"
3996+
3997+
TF = TypeVar("TF", bound=Foo)
3998+
def outer(cls: Type[TF]) -> TF:
3999+
reveal_type(test(cls)) # N: Revealed type is "TF`-1"
4000+
return cls()

test-data/unit/pythoneval.test

Lines changed: 19 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1858,3 +1858,22 @@ _testTupleWithDifferentArgsPy310.py:20: note: Revealed type is "builtins.list[Tu
18581858
_testTupleWithDifferentArgsPy310.py:26: error: Invalid type: try using Literal[1] instead?
18591859
_testTupleWithDifferentArgsPy310.py:27: error: Unexpected "..."
18601860
_testTupleWithDifferentArgsPy310.py:30: note: Revealed type is "builtins.tuple[builtins.object, ...]"
1861+
1862+
[case testEnumIterMetaInference]
1863+
import socket
1864+
from enum import Enum
1865+
from typing import Iterable, Iterator, Type, TypeVar
1866+
1867+
_E = TypeVar("_E", bound=Enum)
1868+
1869+
def enum_iter(cls: Type[_E]) -> Iterable[_E]:
1870+
reveal_type(iter(cls))
1871+
reveal_type(next(iter(cls)))
1872+
return iter(cls)
1873+
1874+
for value in enum_iter(socket.SocketKind):
1875+
reveal_type(value)
1876+
[out]
1877+
_testEnumIterMetaInference.py:8: note: Revealed type is "typing.Iterator[_E`-1]"
1878+
_testEnumIterMetaInference.py:9: note: Revealed type is "_E`-1"
1879+
_testEnumIterMetaInference.py:13: note: Revealed type is "socket.SocketKind"

0 commit comments

Comments
 (0)