[mypyc] Support iterating over a TypedDict

mypyc/irbuild/builder.py

@@ -51,6 +51,7 @@
     ProperType,
     TupleType,
     Type,
+    TypedDictType,
     TypeOfAny,
     UninhabitedType,
     UnionType,
@@ -892,8 +893,12 @@ def get_dict_base_type(self, expr: Expression) -> list[Instance]:
 
         dict_types = []
         for t in types:
-            assert isinstance(t, Instance), t
-            dict_base = next(base for base in t.type.mro if base.fullname == "builtins.dict")
+            if isinstance(t, TypedDictType):
+                t = t.fallback
+                dict_base = next(base for base in t.type.mro if base.fullname == "typing.Mapping")
+            else:
+                assert isinstance(t, Instance), t
+                dict_base = next(base for base in t.type.mro if base.fullname == "builtins.dict")
             dict_types.append(map_instance_to_supertype(t, dict_base))
         return dict_types
 

mypyc/test-data/irbuild-dict.test

@@ -219,6 +219,12 @@ L0:
 
 [case testDictIterationMethods]
 from typing import Dict, Union
+from typing_extensions import TypedDict
+
+class Person(TypedDict):
+    name: str
+    age: int
+
 def print_dict_methods(d1: Dict[int, int], d2: Dict[int, int]) -> None:
     for v in d1.values():
         if v in d2:
@@ -229,6 +235,10 @@ def union_of_dicts(d: Union[Dict[str, int], Dict[str, str]]) -> None:
     new = {}
     for k, v in d.items():
         new[k] = int(v)
+def typeddict(d: Person) -> None:
+    for k, v in d.items():
+        if k == "name":
+            name = v
 [out]
 def print_dict_methods(d1, d2):
     d1, d2 :: dict
@@ -370,6 +380,58 @@ L4:
     r19 = CPy_NoErrOccured()
 L5:
     return 1
+def typeddict(d):
+    d :: dict
+    r0 :: short_int
+    r1 :: native_int
+    r2 :: short_int
+    r3 :: object
+    r4 :: tuple[bool, short_int, object, object]
+    r5 :: short_int
+    r6 :: bool
+    r7, r8 :: object
+    r9 :: str
+    k, v :: object
+    r10 :: str
+    r11 :: object
+    r12 :: int32
+    r13 :: bit
+    r14 :: bool
+    name :: object
+    r15, r16 :: bit
+L0:
+    r0 = 0
+    r1 = PyDict_Size(d)
+    r2 = r1 << 1
+    r3 = CPyDict_GetItemsIter(d)
+L1:
+    r4 = CPyDict_NextItem(r3, r0)
+    r5 = r4[1]
+    r0 = r5
+    r6 = r4[0]
+    if r6 goto L2 else goto L6 :: bool
+L2:
+    r7 = r4[2]
+    r8 = r4[3]
+    r9 = cast(str, r7)
+    k = r9
+    v = r8
+    r10 = 'name'
+    r11 = PyObject_RichCompare(k, r10, 2)
+    r12 = PyObject_IsTrue(r11)
+    r13 = r12 >= 0 :: signed
+    r14 = truncate r12: int32 to builtins.bool
+    if r14 goto L3 else goto L4 :: bool
+L3:
+    name = v
+L4:
+L5:
+    r15 = CPyDict_CheckSize(d, r2)
+    goto L1
+L6:
+    r16 = CPy_NoErrOccured()
+L7:
+    return 1
 
 [case testDictLoadAddress]
 def f() -> None:

test-data/unit/lib-stub/typing_extensions.pyi

@@ -1,5 +1,5 @@
 import typing
-from typing import Any, Mapping, Iterator, NoReturn as NoReturn, Dict, Type
+from typing import Any, Mapping, Iterable, Iterator, NoReturn as NoReturn, Dict, Tuple, Type
 from typing import TYPE_CHECKING as TYPE_CHECKING
 from typing import NewType as NewType, overload as overload
 
@@ -50,6 +50,9 @@ class _TypedDict(Mapping[str, object]):
     # Mypy expects that 'default' has a type variable type.
     def pop(self, k: NoReturn, default: _T = ...) -> object: ...
     def update(self: _T, __m: _T) -> None: ...
+    def items(self) -> Iterable[str, object]: ...
+    def keys(self) -> Iterable[str, object]: ...
+    def values(self) -> Iterable[Tuple[str, object]]: ...
     if sys.version_info < (3, 0):
         def has_key(self, k: str) -> bool: ...
     def __delitem__(self, k: NoReturn) -> None: ...