From 8bfa4f33fc1193941a6bf5c41474f06798b91241 Mon Sep 17 00:00:00 2001 From: Nicolas Dandrimont Date: Wed, 15 Jan 2014 23:55:33 +0100 Subject: [PATCH] Add set comprehensions, dict comprehensions and generator expressions Closes: #14 (woo, two-digit tickets) --- hy/_compat.py | 1 + hy/compiler.py | 113 +++++++++++++++++++++++++++------ tests/native_tests/language.hy | 37 ++++++++++- 3 files changed, 130 insertions(+), 21 deletions(-) diff --git a/hy/_compat.py b/hy/_compat.py index 37bb023..096282a 100644 --- a/hy/_compat.py +++ b/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) diff --git a/hy/compiler.py b/hy/compiler.py index 6e65b30..e9be303 100644 --- a/hy/compiler.py +++ b/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() - - 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) + gen_res, gen = self._compile_generator_iterables(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): diff --git a/tests/native_tests/language.hy b/tests/native_tests/language.hy index 11bf01a..cb75e66 100644 --- a/tests/native_tests/language.hy +++ b/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 [])