Restrict comprehensions to single for and no if clauses.

artiq/compiler/transforms/asttyped_rewriter.py

@@ -57,6 +57,25 @@ class LocalExtractor(algorithm.Visitor):
         for if_ in node.ifs:
             self.visit(node.ifs)
 
+    def visit_generator(self, node):
+        if self.in_root:
+            return
+        self.in_root = True
+        self.visit(list(reversed(node.generators)))
+        self.visit(node.elt)
+
+    visit_ListComp     = visit_generator
+    visit_SetComp      = visit_generator
+    visit_GeneratorExp = visit_generator
+
+    def visit_DictComp(self, node):
+        if self.in_root:
+            return
+        self.in_root = True
+        self.visit(list(reversed(node.generators)))
+        self.visit(node.key)
+        self.visit(node.value)
+
     def visit_root(self, node):
         if self.in_root:
             return
@@ -66,10 +85,6 @@ class LocalExtractor(algorithm.Visitor):
     visit_Module       = visit_root # don't look at inner scopes
     visit_ClassDef     = visit_root
     visit_Lambda       = visit_root
-    visit_DictComp     = visit_root
-    visit_ListComp     = visit_root
-    visit_SetComp      = visit_root
-    visit_GeneratorExp = visit_root
 
     def visit_FunctionDef(self, node):
         if self.in_root:

artiq/compiler/transforms/inferencer.py

@@ -390,11 +390,23 @@ class Inferencer(algorithm.Visitor):
                     node.loc, None)
 
     def visit_ListCompT(self, node):
+        if len(node.generators) > 1:
+            diag = diagnostic.Diagnostic("error",
+                "multiple for clauses in comprehensions are not supported", {},
+                node.generators[1].for_loc)
+            self.engine.process(diag)
+
         self.generic_visit(node)
         self._unify(node.type, builtins.TList(node.elt.type),
                     node.loc, None)
 
     def visit_comprehension(self, node):
+        if any(node.ifs):
+            diag = diagnostic.Diagnostic("error",
+                "if clauses in comprehensions are not supported", {},
+                node.if_locs[0])
+            self.engine.process(diag)
+
         self.generic_visit(node)
         self._unify_iterable(element=node.target, collection=node.iter)
 

lit-test/compiler/inferencer/error_comprehension.py

@@ -0,0 +1,8 @@
+# RUN: %python -m artiq.compiler.testbench.inferencer +diag %s >%t
+# RUN: OutputCheck %s --file-to-check=%t
+
+# CHECK-L: ${LINE:+1}: error: if clauses in comprehensions are not supported
+[x for x in [] if x]
+
+# CHECK-L: ${LINE:+1}: error: multiple for clauses in comprehensions are not supported
+[(x, y) for x in [] for y in []]