Repository: python/mypy
Pull request: [mypyc] Support iterating over a TypedDict
An optimization to make iterating over dict.keys(), dict.values() and dict.items() faster caused mypyc to crash while compiling a TypedDict. This commit fixes `Builder.get_dict_base_type` to properly handle `TypedDictType`.

Fixes mypyc/mypyc#869.

Problem statement:
Using TypedDict in function triggers AssertionError during compilation
```py
from typing import TypedDict

class Movie(TypedDict):
    name: str
    year: int

def test(dictionary: Movie) -> None:
    for key, value in dictionary.items():
        print(key, value)
```

with mypy 0.900 and python 3.9
```sh
$ mypyc typed_dict.py
Traceback (most recent call last):
  File "mypyc/irbuild/builder.py", line 169, in accept
  File "mypy/nodes.py", line 1140, in accept
  File "mypyc/irbuild/visitor.py", line 148, in visit_for_stmt
  File "mypyc/irbuild/statement.py", line 252, in transform_for_stmt
  File "mypyc/irbuild/for_helpers.py", line 60, in for_loop_helper
  File "mypyc/irbuild/for_helpers.py", line 387, in make_for_loop_generator
  File "mypyc/irbuild/builder.py", line 719, in get_dict_item_type
  File "mypyc/irbuild/builder.py", line 711, in get_dict_key_type
  File "mypyc/irbuild/builder.py", line 705, in get_dict_base_type
typed_dict.py:8: AssertionError: 
```

Hints from issue discussion:
Hello, I can confirm this error is still present on

```
Python 3.9.7
mypy 0.920+dev.d263a9698fb59d7a1ae0c7817ef53bb272c8d875
```

Here is the full traceback:

```
Traceback (most recent call last):
  File "build_dir/py_plus_c/build/setup.py", line 5, in <module>
    ext_modules=mypycify(['typed_dict.py'], opt_level="3"),
  File "build_dir/mypy/mypyc/build.py", line 485, in mypycify
    groups, group_cfilenames = mypyc_build(
  File "build_dir/mypy/mypyc/build.py", line 398, in mypyc_build
    group_cfiles, ops_text = generate_c(all_sources, options, groups, fscache,
  File "build_dir/mypy/mypyc/build.py", line 200, in generate_c
    modules, ctext = emitmodule.compile_modules_to_c(
  File "build_dir/mypy/mypyc/codegen/emitmodule.py", line 414, in compile_modules_to_c
    modules = compile_modules_to_ir(result, mapper, compiler_options, errors)
  File "build_dir/mypy/mypyc/codegen/emitmodule.py", line 258, in compile_modules_to_ir
    scc_ir = compile_scc_to_ir(trees, result, mapper, compiler_options, errors)
  File "build_dir/mypy/mypyc/codegen/emitmodule.py", line 209, in compile_scc_to_ir
    modules = build_ir(
  File "/usr/local/lib/python3.9/contextlib.py", line 79, in inner
    return func(*args, **kwds)
  File "build_dir/mypy/mypyc/irbuild/main.py", line 82, in build_ir
    transform_mypy_file(builder, module)
  File "build_dir/mypy/mypyc/irbuild/main.py", line 125, in transform_mypy_file
    builder.accept(node)
  File "build_dir/mypy/mypyc/irbuild/builder.py", line 175, in accept
    node.accept(self.visitor)
  File "build_dir/mypy/mypy/nodes.py", line 696, in accept
    return visitor.visit_func_def(self)
  File "build_dir/mypy/mypyc/irbuild/visitor.py", line 116, in visit_func_def
    transform_func_def(self.builder, fdef)
  File "build_dir/mypy/mypyc/irbuild/function.py", line 69, in transform_func_def
    func_ir, func_reg = gen_func_item(builder, fdef, fdef.name, builder.mapper.fdef_to_sig(fdef))
  File "build_dir/mypy/mypyc/irbuild/function.py", line 306, in gen_func_item
    builder.accept(fitem.body)
  File "build_dir/mypy/mypyc/irbuild/builder.py", line 175, in accept
    node.accept(self.visitor)
  File "build_dir/mypy/mypy/nodes.py", line 1024, in accept
    return visitor.visit_block(self)
  File "build_dir/mypy/mypyc/irbuild/visitor.py", line 125, in visit_block
    transform_block(self.builder, block)
  File "build_dir/mypy/mypyc/irbuild/statement.py", line 46, in transform_block
    builder.accept(stmt)
  File "build_dir/mypy/mypyc/irbuild/builder.py", line 175, in accept
    node.accept(self.visitor)
  File "build_dir/mypy/mypy/nodes.py", line 1150, in accept
    return visitor.visit_for_stmt(self)
  File "build_dir/mypy/mypyc/irbuild/visitor.py", line 148, in visit_for_stmt
    transform_for_stmt(self.builder, stmt)
  File "build_dir/mypy/mypyc/irbuild/statement.py", line 254, in transform_for_stmt
    for_loop_helper(builder, s.index, s.expr, body,
  File "build_dir/mypy/mypyc/irbuild/for_helpers.py", line 60, in for_loop_helper
    for_gen = make_for_loop_generator(builder, index, expr, body_block, normal_loop_exit, line)
  File "build_dir/mypy/mypyc/irbuild/for_helpers.py", line 389, in make_for_loop_generator
    target_type = builder.get_dict_item_type(expr.callee.expr)
  File "build_dir/mypy/mypyc/irbuild/builder.py", line 761, in get_dict_item_type
    key_type = self.get_dict_key_type(expr)
  File "build_dir/mypy/mypyc/irbuild/builder.py", line 753, in get_dict_key_type
    dict_base_type = self.get_dict_base_type(expr)
  File "build_dir/mypy/mypyc/irbuild/builder.py", line 747, in get_dict_base_type
    assert isinstance(target_type, Instance)
typed_dict.py:8: AssertionError: 
```

I have also noted that wrapping the `dictionary` variable with `dict` makes the compilation successful:

```
from typing import TypedDict

class Movie(TypedDict):
    name: str
    year: int

def test(dictionary: Movie) -> None:
    for key, value in dict(dictionary).items():
        print(key, value)
```

Commit message: [mypyc] Support iterating over a TypedDict

An optimization to make iterating over dict.keys(), dict.values() and
dict.items() faster caused mypyc to crash while compiling a TypedDict.
This commit fixes `Builder.get_dict_base_type` to properly handle
TypedDictType.

Patch to review:
diff --git a/mypyc/irbuild/builder.py b/mypyc/irbuild/builder.py
index f37fae608083..e001808a5c78 100644
--- a/mypyc/irbuild/builder.py
+++ b/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
 
diff --git a/mypyc/test-data/irbuild-dict.test b/mypyc/test-data/irbuild-dict.test
index 99643b9451f0..a51a5832a5fd 100644
--- a/mypyc/test-data/irbuild-dict.test
+++ b/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:
diff --git a/test-data/unit/lib-stub/typing_extensions.pyi b/test-data/unit/lib-stub/typing_extensions.pyi
index 22b895971521..5dbde18bab63 100644
--- a/test-data/unit/lib-stub/typing_extensions.pyi
+++ b/test-data/unit/lib-stub/typing_extensions.pyi
@@ -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) -> dict_items[str, object]: ...
+    def keys(self) -> dict_keys[str, object]: ...
+    def values(self) -> dict_values[str, object]: ...
     if sys.version_info < (3, 0):
         def has_key(self, k: str) -> bool: ...
     def __delitem__(self, k: NoReturn) -> None: ...


Code review comments:
diff_hunk: @@ -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) -> dict_items[str, object]: ...
+    def keys(self) -> dict_keys[str, object]: ...
+    def values(self) -> dict_values[str, object]: ...
path: test-data/unit/lib-stub/typing_extensions.pyi
text: @user1:
The return types aren't imported anywhere. For some reason we don't generate errors even though these are undefined. Maybe we are ignoring errors in test stubs? 

You should be able to use a less specific type. These are adapted from mypyc test stubs:
```py
    def keys(self) -> Iterable[str]: pass
    def values(self) -> Iterable[object]: pass
    def items(self) -> Iterable[Tuple[str, object]]: pass
```

@author:
Ah, good catch! And yeah I have no idea what's up with the lack of errors. I opened an issue to note that we should investigate this later: https://github.com/mypyc/mypyc/issues/976.