diff --git a/artiq/compiler/transforms/asttyped_rewriter.py b/artiq/compiler/transforms/asttyped_rewriter.py index de72fce3e..e89e7ac9c 100644 --- a/artiq/compiler/transforms/asttyped_rewriter.py +++ b/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: diff --git a/artiq/compiler/transforms/inferencer.py b/artiq/compiler/transforms/inferencer.py index 7fbf09813..dbb9aeb04 100644 --- a/artiq/compiler/transforms/inferencer.py +++ b/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) diff --git a/lit-test/compiler/inferencer/error_comprehension.py b/lit-test/compiler/inferencer/error_comprehension.py new file mode 100644 index 000000000..d586dd657 --- /dev/null +++ b/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 []]