added list comprehension

Author: craigbarratt

custom_components/pyscript/eval.py

@@ -981,6 +981,26 @@ async def ast_list(self, arg):
         if isinstance(arg.ctx, ast.Load):
             return await self.eval_elt_list(arg.elts)
 
+    async def listcomp_loop(self, generators, elt):
+        """Recursive list 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.append(await self.aeval(elt))
+                else:
+                    out += await self.listcomp_loop(generators[1:], elt)
+        return out
+
+    async def ast_listcomp(self, arg):
+        """Evaluate list comprehension."""
+        return await self.listcomp_loop(arg.generators, arg.elt)
+
     async def ast_tuple(self, arg):
         """Evaluate Tuple."""
         return tuple(await self.eval_elt_list(arg.elts))

tests/custom_components/pyscript/test_unit_eval.py

@@ -178,6 +178,33 @@
     ["import random as rand, math as m\n[rand.uniform(10,10), m.sqrt(1024)]", [10, 32]],
     ["import cmath\ncmath.sqrt(complex(3, 4))", 2 + 1j],
     ["from math import sqrt as sqroot\nsqroot(1024)", 32],
+    ["[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}