diff --git a/hy/compiler.py b/hy/compiler.py index 9c08cff..92c611a 100644 --- a/hy/compiler.py +++ b/hy/compiler.py @@ -349,6 +349,32 @@ class HyASTCompiler(object): col_offset=expr.start_column, ctx=ast.Load()) + @builds("list_comp") + def compile_list_comprehension(self, expr): + # (list-comp expr (target iter)) + expr.pop(0) + thing = self.compile(expr.pop(0)) + ident, gen = expr.pop(0) + + ident = self.compile(ident) + gen = self.compile(gen) + + if isinstance(ident, ast.Tuple): + for x in ident.elts: + x.ctx = ast.Store() + ident.ctx = ast.Store() + + return ast.ListComp( + lineno=expr.start_line, + col_offset=expr.start_column, + elt=thing, + generators=[ + ast.comprehension( + target=ident, + iter=gen, + ifs=[self.compile(x) for x in expr]) + ]) + @builds("kwapply") def compile_kwapply_expression(self, expr): expr.pop(0) # kwapply diff --git a/tests/native_tests/language.hy b/tests/native_tests/language.hy index 0cda171..62291f9 100644 --- a/tests/native_tests/language.hy +++ b/tests/native_tests/language.hy @@ -200,3 +200,10 @@ "NATIVE: test with" (with-as (open "README.md" "r") fd (pass))) + + +(defn test-comprehensions [] + "NATIVE: test list comprehensions" + (assert (= (list-comp (* x 2) (x (range 2))) [0 2])) + (assert (= (list-comp (* x 2) (x (range 4)) (% x 2)) [2 6])) + (assert (= (list-comp (* y 2) ((, x y) (.iteritems {"1" 1 "2" 2}))) [2 4])))