added dict comprehension

Author: craigbarratt

custom_components/pyscript/eval.py

@@ -1016,11 +1016,32 @@ async def ast_dict(self, arg):
                 val[await self.aeval(key_ast)] = this_val
         return val
 
+    async def dictcomp_loop(self, generators, key, value):
+        """Recursive dict comprehension."""
+        out = {}
+        gen = generators[0]
+        for loop_var in await self.aeval(gen.iter):
+            await self.recurse_assign(gen.target, loop_var)
+            for cond in gen.ifs:
+                if not await self.aeval(cond):
+                    break
+            else:
+                if len(generators) == 1:
+                    out[await self.aeval(key)] = await self.aeval(value)
+                else:
+                    out.update(await self.dictcomp_loop(generators[1:], key, value))
+        return out
+
+    async def ast_dictcomp(self, arg):
+        """Evaluate dict comprehension."""
+        return await self.dictcomp_loop(arg.generators, arg.key, arg.value)
+
     async def ast_set(self, arg):
         """Evaluate set."""
-        return {
-            elt for elt in await self.eval_elt_list(arg.elts)
-        }  # pylint: disable=unnecessary-comprehension
+        ret = set()
+        for elt in await self.eval_elt_list(arg.elts):
+            ret.add(elt)
+        return ret
 
     async def setcomp_loop(self, generators, elt):
         """Recursive list comprehension."""

tests/custom_components/pyscript/test_unit_eval.py

@@ -232,6 +232,29 @@ def no_op(i):
 """,
         [{0, 1, 2, 4, 6}, 13],
     ],
+    ["{str(i):i for i in range(5) if i != 3}", {"0": 0, "1": 1, "2": 2, "4": 4}],
+    [
+        "{v:k for k,v in {str(i):i for i in range(5)}.items()}",
+        {0: "0", 1: "1", 2: "2", 3: "3", 4: "4"},
+    ],
+    [
+        "{f'{i}+{k}':i+k for i in range(3) for k in range(3)}",
+        {
+            "0+0": 0,
+            "0+1": 1,
+            "0+2": 2,
+            "1+0": 1,
+            "1+1": 2,
+            "1+2": 3,
+            "2+0": 2,
+            "2+1": 3,
+            "2+2": 4,
+        },
+    ],
+    [
+        "{f'{i}+{k}':i+k for i in range(3) for k in range(3) if k <= i}",
+        {"0+0": 0, "1+0": 1, "1+1": 2, "2+0": 2, "2+1": 3, "2+2": 4},
+    ],
     [
         """
 d = {"x": 1, "y": 2, "z": 3}