Add set comprehensions, dict comprehensions and generator expressions

Closes: #14 (woo, two-digit tickets)

hy/_compat.py

@@ -38,6 +38,7 @@ except ImportError:
                        (x >> 24) & 0xff]))
 import sys
 
+PY27 = sys.version_info >= (2, 7)
 PY3 = sys.version_info[0] >= 3
 PY33 = sys.version_info >= (3, 3)
 PY34 = sys.version_info >= (3, 4)

hy/compiler.py

@@ -39,7 +39,7 @@ from hy.errors import HyCompileError, HyTypeError
 
 import hy.macros
 from hy.macros import require, macroexpand
-from hy._compat import str_type, long_type, PY33, PY3, PY34
+from hy._compat import str_type, long_type, PY27, PY33, PY3, PY34
 import hy.importer
 
 import traceback
@@ -1265,41 +1265,114 @@ class HyASTCompiler(object):
                          ctx=ast.Load())
         return ret
 
+    def _compile_generator_iterables(self, trailers):
+        """Helper to compile the "trailing" parts of comprehensions:
+        generators and conditions"""
+
+        generators = trailers.pop(0)
+
+        cond = self.compile(trailers.pop(0)) if trailers != [] else Result()
+
+        gen_it = iter(generators)
+        paired_gens = zip(gen_it, gen_it)
+
+        gen_res = Result()
+        gen = []
+        for target, iterable in paired_gens:
+            comp_target = self.compile(target)
+            target = self._storeize(comp_target)
+            gen_res += self.compile(iterable)
+            gen.append(ast.comprehension(
+                target=target,
+                iter=gen_res.force_expr,
+                ifs=[]))
+
+        if cond.expr:
+            gen[-1].ifs.append(cond.expr)
+
+        return gen_res + cond, gen
+
     @builds("list_comp")
     @checkargs(min=2, max=3)
     def compile_list_comprehension(self, expr):
         # (list-comp expr (target iter) cond?)
         expr.pop(0)
         expression = expr.pop(0)
-        tar_it = iter(expr.pop(0))
-        targets = zip(tar_it, tar_it)
 
-        cond = self.compile(expr.pop(0)) if expr != [] else Result()
+        gen_res, gen = self._compile_generator_iterables(expr)
-
-        generator_res = Result()
-        generators = []
-        for target, iterable in targets:
-            comp_target = self.compile(target)
-            target = self._storeize(comp_target)
-            generator_res += self.compile(iterable)
-            generators.append(ast.comprehension(
-                target=target,
-                iter=generator_res.force_expr,
-                ifs=[]))
-
-        if cond.expr:
-            generators[-1].ifs.append(cond.expr)
 
         compiled_expression = self.compile(expression)
-        ret = compiled_expression + generator_res + cond
+        ret = compiled_expression + gen_res
         ret += ast.ListComp(
             lineno=expr.start_line,
             col_offset=expr.start_column,
             elt=compiled_expression.force_expr,
-            generators=generators)
+            generators=gen)
 
         return ret
 
+    @builds("set_comp")
+    @checkargs(min=2, max=3)
+    def compile_set_comprehension(self, expr):
+        if PY27:
+            ret = self.compile_list_comprehension(expr)
+            expr = ret.expr
+            ret.expr = ast.SetComp(
+                lineno=expr.lineno,
+                col_offset=expr.col_offset,
+                elt=expr.elt,
+                generators=expr.generators)
+
+            return ret
+
+        expr[0] = HySymbol("list_comp").replace(expr[0])
+        expr = HyExpression([HySymbol("set"), expr]).replace(expr)
+        return self.compile(expr)
+
+    @builds("dict_comp")
+    @checkargs(min=3, max=4)
+    def compile_dict_comprehension(self, expr):
+        if PY27:
+            expr.pop(0)  # dict-comp
+            key = expr.pop(0)
+            value = expr.pop(0)
+
+            gen_res, gen = self._compile_generator_iterables(expr)
+
+            compiled_key = self.compile(key)
+            compiled_value = self.compile(value)
+            ret = compiled_key + compiled_value + gen_res
+            ret += ast.DictComp(
+                lineno=expr.start_line,
+                col_offset=expr.start_column,
+                key=compiled_key.force_expr,
+                value=compiled_value.force_expr,
+                generators=gen)
+
+            return ret
+
+        # In Python 2.6, turn (dict-comp key value [foo]) into
+        # (dict (list-comp (, key value) [foo]))
+
+        expr[0] = HySymbol("list_comp").replace(expr[0])
+        expr[1:3] = [HyExpression(
+            [HySymbol(",")] +
+            expr[1:3]
+        ).replace(expr[1])]
+        expr = HyExpression([HySymbol("dict"), expr]).replace(expr)
+        return self.compile(expr)
+
+    @builds("genexpr")
+    def compile_genexpr(self, expr):
+        ret = self.compile_list_comprehension(expr)
+        expr = ret.expr
+        ret.expr = ast.GeneratorExp(
+            lineno=expr.lineno,
+            col_offset=expr.col_offset,
+            elt=expr.elt,
+            generators=expr.generators)
+        return ret
+
     @builds("apply")
     @checkargs(min=1, max=3)
     def compile_apply_expression(self, expr):

tests/native_tests/language.hy

@@ -525,7 +525,7 @@
     (assert (= x 3))))
 
 
-(defn test-comprehensions []
+(defn test-list-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]))
@@ -536,6 +536,41 @@
   (assert (= (list-comp j (j [1 2])) [1 2])))
 
 
+(defn test-set-comprehensions []
+  "NATIVE: test set comprehensions"
+  (assert (instance? set (set-comp x [x (range 2)])))
+  (assert (= (set-comp (* x 2) (x (range 2))) (set [0 2])))
+  (assert (= (set-comp (* x 2) (x (range 4)) (% x 2)) (set [2 6])))
+  (assert (= (set-comp (* y 2) ((, x y) (.items {"1" 1 "2" 2})))
+             (set [2 4])))
+  (assert (= (set-comp (, x y) (x (range 2) y (range 2)))
+             (set [(, 0 0) (, 0 1) (, 1 0) (, 1 1)])))
+  (assert (= (set-comp j (j [1 2])) (set [1 2]))))
+
+
+(defn test-dict-comprehensions []
+  "NATIVE: test dict comprehensions"
+  (assert (instance? dict (dict-comp x x [x (range 2)])))
+  (assert (= (dict-comp x (* x 2) (x (range 2))) {1 2 0 0}))
+  (assert (= (dict-comp x (* x 2) (x (range 4)) (% x 2)) {3 6 1 2}))
+  (assert (= (dict-comp x (* y 2) ((, x y) (.items {"1" 1 "2" 2})))
+             {"2" 4 "1" 2}))
+  (assert (= (dict-comp (, x y) (+ x y) (x (range 2) y (range 2)))
+             {(, 0 0) 0 (, 1 0) 1 (, 0 1) 1 (, 1 1) 2})))
+
+
+(defn test-generator-expressions []
+  "NATIVE: test generator expressions"
+  (assert (not (instance? list (genexpr x [x (range 2)]))))
+  (assert (= (list (genexpr (* x 2) (x (range 2)))) [0 2]))
+  (assert (= (list (genexpr (* x 2) (x (range 4)) (% x 2))) [2 6]))
+  (assert (= (list (sorted (genexpr (* y 2) ((, x y) (.items {"1" 1 "2" 2})))))
+             [2 4]))
+  (assert (= (list (genexpr (, x y) (x (range 2) y (range 2))))
+             [(, 0 0) (, 0 1) (, 1 0) (, 1 1)]))
+  (assert (= (list (genexpr j (j [1 2]))) [1 2])))
+
+
 (defn test-defn-order []
   "NATIVE: test defn evaluation order"
   (setv acc [])