Merge pull request #448 from olasd/feature/genexpr
Add set comprehensions, dict comprehensions and generator expressions
This commit is contained in:
commit
d34c22eeac
@ -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)
|
||||
|
113
hy/compiler.py
113
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):
|
||||
|
@ -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 [])
|
||||
|
Loading…
Reference in New Issue
Block a user