sklam · Feb 23, 2021 · Feb 19, 2021 · Feb 23, 2021 · esc · Feb 23, 2021
diff --git a/docs/source/reference/pysupported.rst b/docs/source/reference/pysupported.rst
@@ -796,6 +796,23 @@ range of possible failures. However, the dictionary can be safely read from
 multiple threads as long as the contents of the dictionary do not
 change during the parallel access.
 
+Dictionary comprehension
+''''''''''''''''''''''''
+
+Numba supports dictionary comprehension under the assumption that a
+``numba.typed.Dict`` instance can be created from the comprehension.  For
+example::
+
+  In [1]: from numba import njit
+
+  In [2]: @njit
+  ...: def foo(n):
+  ...:     return {i: i**2 for i in range(n)}
+  ...:
-  In [2]: @njit
-  ...: def foo(n):
-  ...:     return {i: i**2 for i in range(n)}
-  ...:
+In [2]: @njit
+   ...: def foo(n):
+   ...:     return {i: i**2 for i in range(n)}
+   ...:
-  In [2]: @njit
-  ...: def foo(n):
-  ...:     return {i: i**2 for i in range(n)}
-  ...:
+In [2]: @njit
+   ...: def foo(n):
+   ...:     return {i: i**2 for i in range(n)}
+   ...:
+
+  In [3]: foo(3)
+  Out[3]: DictType[int64,int64]<iv=None>({0: 0, 1: 1, 2: 4})
+
 .. _feature-dict-initial-value:
 
 Initial Values

diff --git a/numba/core/byteflow.py b/numba/core/byteflow.py
@@ -916,6 +916,20 @@ def op_BUILD_MAP(self, state, inst):
         state.append(inst, items=items[::-1], size=count, res=dct)
         state.push(dct)
 
+    def op_MAP_ADD(self, state, inst):
+        # NOTE: https://docs.python.org/3/library/dis.html#opcode-MAP_ADD
+        # Python >= 3.8: TOS and TOS1 are value and key respectively
+        # Python < 3.8: TOS and TOS1 are key and value respectively
+        TOS = state.pop()
+        TOS1 = state.pop()
+        key, value = (TOS, TOS1) if PYVERSION < (3, 8) else (TOS1, TOS)
+        index = inst.arg
+>
+        setitemvar = state.make_temp()
+        res = state.make_temp()
+        state.append(inst, key=key, value=value,
+                     setitemvar=setitemvar, res=res)
+
     def op_BUILD_SET(self, state, inst):
         count = inst.arg
         # Note: related python bug http://bugs.python.org/issue26020

diff --git a/numba/core/dataflow.py b/numba/core/dataflow.py
@@ -202,6 +202,16 @@ def op_BUILD_MAP(self, info, inst):
         info.append(inst, items=items[::-1], size=count, res=dct)
         info.push(dct)
 
+    def op_MAP_ADD(self, info, inst):
+        key = info.pop()
+        value = info.pop()
+        index = inst.arg
+>
+        setitemvar = info.make_temp()
+        res = info.make_temp()
+        info.append(inst, key=key, value=value,
+                     setitemvar=setitemvar, res=res)
+
     def op_BUILD_SET(self, info, inst):
         count = inst.arg
         # Note: related python bug http://bugs.python.org/issue26020

diff --git a/numba/core/interpreter.py b/numba/core/interpreter.py
@@ -1803,6 +1803,16 @@ def op_LIST_EXTEND(self, inst, target, value, extendvar, res):
                                       loc=self.loc)
             self.store(value=extendinst, name=res)
 
+    def op_MAP_ADD(self, inst, target, key, value, setitemvar, res):
+>
+        key = self.get(key)
+        value = self.get(value)
+        setitemattr = ir.Expr.getattr(target, '__setitem__', loc=self.loc)
+        self.store(value=setitemattr, name=setitemvar)
+        appendinst = ir.Expr.call(self.get(setitemvar), (key, value,), (),
+                                  loc=self.loc)
+        self.store(value=appendinst, name=res)
+
     def op_LOAD_ASSERTION_ERROR(self, inst, res):
         gv_fn = ir.Global("AssertionError", AssertionError, loc=self.loc)
         self.store(value=gv_fn, name=res)

diff --git a/numba/tests/test_dictobject.py b/numba/tests/test_dictobject.py
@@ -1534,6 +1534,40 @@ def foo(k1, k2, v):
             },
         )
 
+    def test_comprehension_basic(self):
+        @njit
+        def foo():
+            return {i: 2 * i for i in range(10)}
+
+        self.assertEqual(foo(), foo.py_func())
+
+    def test_comprehension_basic_mixed_type(self):
+        @njit
+        def foo():
+            return {i: float(j) for i, j in zip(range(10), range(10, 0, -1))}
+
+        self.assertEqual(foo(), foo.py_func())
+
+    def test_comprehension_involved(self):
+        @njit
+        def foo():
+            a = {0: 'A', 1: 'B', 2: 'C'}
+            return {3 + i: a[i] for i in range(3)}
+
+        self.assertEqual(foo(), foo.py_func())
+
+    def test_comprehension_fail_mixed_type(self):
+        @njit
+        def foo():
+            a = {0: 'A', 1: 'B', 2: 1j}
+            return {3 + i: a[i] for i in range(3)}
+
+        with self.assertRaises(TypingError) as e:
+            foo()
+
+        excstr = str(e.exception)
+        self.assertIn("Cannot cast complex128 to unicode_type", excstr)
+
 
 class TestNonCompiledInfer(TestCase):
     def test_check_untyped_dict_ops(self):

diff --git a/numba/typed/dictobject.py b/numba/typed/dictobject.py
@@ -671,6 +671,7 @@ def impl(d):
     return impl
 
 
+@overload_method(types.DictType, '__setitem__')
 @overload(operator.setitem)
 def impl_setitem(d, key, value):
     if not isinstance(d, types.DictType):