Add set comprehensions, dict comprehensions and generator expressions

Closes: #14 (woo, two-digit tickets)
This commit is contained in:
Nicolas Dandrimont 2014-01-15 23:55:33 +01:00
parent da59b20471
commit 8bfa4f33fc
3 changed files with 130 additions and 21 deletions

View File

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

View File

@ -39,7 +39,7 @@ from hy.errors import HyCompileError, HyTypeError
import hy.macros import hy.macros
from hy.macros import require, macroexpand 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 hy.importer
import traceback import traceback
@ -1265,41 +1265,114 @@ class HyASTCompiler(object):
ctx=ast.Load()) ctx=ast.Load())
return ret 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") @builds("list_comp")
@checkargs(min=2, max=3) @checkargs(min=2, max=3)
def compile_list_comprehension(self, expr): def compile_list_comprehension(self, expr):
# (list-comp expr (target iter) cond?) # (list-comp expr (target iter) cond?)
expr.pop(0) expr.pop(0)
expression = 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) compiled_expression = self.compile(expression)
ret = compiled_expression + generator_res + cond ret = compiled_expression + gen_res
ret += ast.ListComp( ret += ast.ListComp(
lineno=expr.start_line, lineno=expr.start_line,
col_offset=expr.start_column, col_offset=expr.start_column,
elt=compiled_expression.force_expr, elt=compiled_expression.force_expr,
generators=generators) generators=gen)
return ret 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") @builds("apply")
@checkargs(min=1, max=3) @checkargs(min=1, max=3)
def compile_apply_expression(self, expr): def compile_apply_expression(self, expr):

View File

@ -525,7 +525,7 @@
(assert (= x 3)))) (assert (= x 3))))
(defn test-comprehensions [] (defn test-list-comprehensions []
"NATIVE: test list comprehensions" "NATIVE: test list comprehensions"
(assert (= (list-comp (* x 2) (x (range 2))) [0 2])) (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 (* x 2) (x (range 4)) (% x 2)) [2 6]))
@ -536,6 +536,41 @@
(assert (= (list-comp j (j [1 2])) [1 2]))) (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 [] (defn test-defn-order []
"NATIVE: test defn evaluation order" "NATIVE: test defn evaluation order"
(setv acc []) (setv acc [])