added set comprehensions

Author: craigbarratt

custom_components/pyscript/eval.py

@@ -1018,7 +1018,29 @@ async def ast_dict(self, arg):
 
     async def ast_set(self, arg):
         """Evaluate set."""
-        return {elt for elt in await self.eval_elt_list(arg.elts)}  # pylint: disable=unnecessary-comprehension
+        return {
+            elt for elt in await self.eval_elt_list(arg.elts)
+        }  # pylint: disable=unnecessary-comprehension
+
+    async def setcomp_loop(self, generators, elt):
+        """Recursive list comprehension."""
+        out = set()
+        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.add(await self.aeval(elt))
+                else:
+                    out.update(await self.setcomp_loop(generators[1:], elt))
+        return out
+
+    async def ast_setcomp(self, arg):
+        """Evaluate set comprehension."""
+        return await self.setcomp_loop(arg.generators, arg.elt)
 
     async def ast_subscript(self, arg):
         """Evaluate subscript."""

tests/custom_components/pyscript/test_unit_eval.py

@@ -205,6 +205,33 @@ def no_op(i):
 """,
         [0, 1, 2, 4, 6, 13],
     ],
+    ["{i for i in range(7) if i != 5 if i != 3}", {0, 1, 2, 4, 6}],
+    [
+        """
+matrix = [[1, 2, 3], [4, 5], [6, 7, 8, 9]]
+{val for sublist in matrix for val in sublist if val != 8}
+""",
+        {1, 2, 3, 4, 5, 6, 7, 9},
+    ],
+    [
+        """
+matrix = [[1, 2, 3], [4, 5], [6, 7, 8, 9]]
+{val for sublist in matrix if sublist[0] != 4 for val in sublist if val != 8}
+""",
+        {1, 2, 3, 6, 7, 9},
+    ],
+    [
+        """
+# check short-circuit of nested if
+cnt = 0
+def no_op(i):
+    global cnt
+    cnt += 1
+    return i
+[{i for i in range(7) if no_op(i) != 5 if no_op(i) != 3}, cnt]
+""",
+        [{0, 1, 2, 4, 6}, 13],
+    ],
     [
         """
 d = {"x": 1, "y": 2, "z": 3}