beewareGH-683: Support multiple if conditions in list comprehensions

vanjaw · Oct 31, 2017 · ac1a5c6 · ac1a5c6
1 parent 5d1b04e
commit ac1a5c6
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 3 deletions.
diff --git a/tests/structures/test_list_comprehension.py b/tests/structures/test_list_comprehension.py
@@ -12,8 +12,8 @@ def test_syntax(self):
 
     def test_list_comprehension_with_if_condition(self):
         self.assertCodeExecution("""
-            x = [1, 2, 3, 4, 5]
-            print([v**2 for v in x if v % 2 == 0])
+            print([v for v in range(100) if v % 2 == 0])
+            print([v for v in range(100) if v % 2 == 0 if v % 3 == 0])
             """)
 
     def test_method(self):

diff --git a/voc/python/ast.py b/voc/python/ast.py
@@ -1244,7 +1244,9 @@ def visit_ListComp(self, node):
                 self.visit(generator.target)
 
         if node.generators[0].ifs:
-            self.visit(node.generators[0].ifs[0])
+            self.visit(
+                ast.BoolOp(ast.And(), node.generators[0].ifs)
+            )
 
             self.context.add_opcodes(
                 IF([python.Object.as_boolean()], JavaOpcodes.IFEQ),