Skip to content

Commit

Permalink
Support protocol inference for Type[T] via metaclass (#14554)
Browse files Browse the repository at this point in the history
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.
  • Loading branch information
ilevkivskyi authored Jan 29, 2023
1 parent c4ecd2b commit 8af3af3
Show file tree
Hide file tree
Showing 3 changed files with 51 additions and 0 deletions.
11 changes: 11 additions & 0 deletions mypy/constraints.py
Original file line number Diff line number Diff line change
Expand Up @@ -619,6 +619,17 @@ def visit_instance(self, template: Instance) -> list[Constraint]:
actual.item, template, subtype, template, class_obj=True
)
)
if self.direction == SUPERTYPE_OF:
# Infer constraints for Type[T] via metaclass of T when it makes sense.
a_item = actual.item
if isinstance(a_item, TypeVarType):
a_item = get_proper_type(a_item.upper_bound)
if isinstance(a_item, Instance) and a_item.type.metaclass_type:
res.extend(
self.infer_constraints_from_protocol_members(
a_item.type.metaclass_type, template, actual, template
)
)

if isinstance(actual, Overloaded) and actual.fallback is not None:
actual = actual.fallback
Expand Down
21 changes: 21 additions & 0 deletions test-data/unit/check-protocols.test
Original file line number Diff line number Diff line change
Expand Up @@ -3977,3 +3977,24 @@ class C:
DEFAULT: ClassVar[C]

x: P = C()

[case testInferenceViaTypeTypeMetaclass]
from typing import Iterator, Iterable, TypeVar, Type

M = TypeVar("M")

class Meta(type):
def __iter__(self: Type[M]) -> Iterator[M]: ...
class Foo(metaclass=Meta): ...

T = TypeVar("T")
def test(x: Iterable[T]) -> T: ...

reveal_type(test(Foo)) # N: Revealed type is "__main__.Foo"
t_foo: Type[Foo]
reveal_type(test(t_foo)) # N: Revealed type is "__main__.Foo"

TF = TypeVar("TF", bound=Foo)
def outer(cls: Type[TF]) -> TF:
reveal_type(test(cls)) # N: Revealed type is "TF`-1"
return cls()
19 changes: 19 additions & 0 deletions test-data/unit/pythoneval.test
Original file line number Diff line number Diff line change
Expand Up @@ -1858,3 +1858,22 @@ _testTupleWithDifferentArgsPy310.py:20: note: Revealed type is "builtins.list[Tu
_testTupleWithDifferentArgsPy310.py:26: error: Invalid type: try using Literal[1] instead?
_testTupleWithDifferentArgsPy310.py:27: error: Unexpected "..."
_testTupleWithDifferentArgsPy310.py:30: note: Revealed type is "builtins.tuple[builtins.object, ...]"

[case testEnumIterMetaInference]
import socket
from enum import Enum
from typing import Iterable, Iterator, Type, TypeVar

_E = TypeVar("_E", bound=Enum)

def enum_iter(cls: Type[_E]) -> Iterable[_E]:
reveal_type(iter(cls))
reveal_type(next(iter(cls)))
return iter(cls)

for value in enum_iter(socket.SocketKind):
reveal_type(value)
[out]
_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"

0 comments on commit 8af3af3

Please sign in to comment.