Skip to content

Commit 3078552

Browse files
ncwncw
committed
Generator, List, Set and Dictionary comprehensions
1 parent 817f25d commit 3078552

File tree

3 files changed

+97
-12
lines changed

3 files changed

+97
-12
lines changed

‎parser/grammar.y

+35-11
Original file line numberDiff line numberDiff line change
@@ -104,7 +104,7 @@ func tupleOrExpr(pos ast.Pos, elts []ast.Expr, optional_comma bool) ast.Expr {
104104
%type <stmt> compound_stmt small_stmt expr_stmt del_stmt pass_stmt flow_stmt import_stmt global_stmt nonlocal_stmt assert_stmt break_stmt continue_stmt return_stmt raise_stmt yield_stmt
105105
%type <op> augassign
106106
%type <expr> expr_or_star_expr expr star_expr xor_expr and_expr shift_expr arith_expr term factor power trailer atom test_or_star_expr test not_test lambdef test_nocond lambdef_nocond or_test and_test comparison testlist testlist_star_expr yield_expr_or_testlist yield_expr yield_expr_or_testlist_star_expr dictorsetmaker
107-
%type <exprs> exprlist testlistraw
107+
%type <exprs> exprlist testlistraw comp_if comp_iter
108108
%type <exprsStack> expr_or_star_exprs test_or_star_exprs tests test_colon_tests
109109
%type <trailers> trailers
110110
%type <cmpop> comp_op
@@ -1103,9 +1103,7 @@ atom:
11031103
}
11041104
| '(' test_or_star_expr comp_for ')'
11051105
{
1106-
// FIXME
1107-
panic("not implemented")
1108-
$$ = nil
1106+
$$ = &ast.GeneratorExp{ExprBase: ast.ExprBase{$<pos>$}, Elt: $2, Generators: $3}
11091107
}
11101108
| '(' test_or_star_exprs optional_comma ')'
11111109
{
@@ -1117,9 +1115,7 @@ atom:
11171115
}
11181116
| '[' test_or_star_expr comp_for ']'
11191117
{
1120-
// FIXME
1121-
panic("not implemented")
1122-
$$ = nil
1118+
$$ = &ast.ListComp{ExprBase: ast.ExprBase{$<pos>$}, Elt: $2, Generators: $3}
11231119
}
11241120
| '[' test_or_star_exprs optional_comma ']'
11251121
{
@@ -1244,6 +1240,7 @@ exprlist:
12441240
expr_or_star_exprs optional_comma
12451241
{
12461242
$$ = $1.Pop()
1243+
$<comma>$ = $2
12471244
}
12481245

12491246
testlist:
@@ -1329,23 +1326,50 @@ argument:
13291326

13301327
comp_iter:
13311328
comp_for
1329+
{
1330+
$<comprehensions>$ = $1
1331+
$$ = nil
1332+
}
13321333
| comp_if
1334+
{
1335+
$<comprehensions>$ = $<comprehensions>1
1336+
$$ = $1
1337+
}
13331338

13341339
comp_for:
13351340
FOR exprlist IN or_test
13361341
{
1337-
// FIXME
1338-
$$ = nil
1342+
c := ast.Comprehension{
1343+
Target: tupleOrExpr($<pos>$, $2, $<comma>2),
1344+
Iter: $4,
1345+
}
1346+
c.Target.(ast.SetCtxer).SetCtx(ast.Store)
1347+
$$ = []ast.Comprehension{c}
13391348
}
13401349
| FOR exprlist IN or_test comp_iter
13411350
{
1342-
// FIXME
1343-
$$ = nil
1351+
c := ast.Comprehension{
1352+
Target: tupleOrExpr($<pos>$, $2, $<comma>2),
1353+
Iter: $4,
1354+
Ifs: $5,
1355+
}
1356+
c.Target.(ast.SetCtxer).SetCtx(ast.Store)
1357+
$$ = []ast.Comprehension{c}
1358+
$$ = append($$, $<comprehensions>5...)
13441359
}
13451360

13461361
comp_if:
13471362
IF test_nocond
1363+
{
1364+
$$ = []ast.Expr{$2}
1365+
$<comprehensions>$ = nil
1366+
}
13481367
| IF test_nocond comp_iter
1368+
{
1369+
$$ = []ast.Expr{$2}
1370+
$$ = append($$, $3...)
1371+
$<comprehensions>$ = $<comprehensions>3
1372+
}
13491373

13501374
// not used in grammar, but may appear in "node" passed from Parser to Compiler
13511375
// encoding_decl: NAME

‎parser/grammar_test.go

+29-1
Original file line numberDiff line numberDiff line change
@@ -59,6 +59,34 @@ func TestGrammar(t *testing.T) {
5959
{"[1,]", "eval", "Expression(body=List(elts=[Num(n=1)], ctx=Load()))"},
6060
{"[1,2]", "eval", "Expression(body=List(elts=[Num(n=1), Num(n=2)], ctx=Load()))"},
6161
{"[1,2,]", "eval", "Expression(body=List(elts=[Num(n=1), Num(n=2)], ctx=Load()))"},
62+
{"( a for a in ab )", "eval", "Expression(body=GeneratorExp(elt=Name(id='a', ctx=Load()), generators=[comprehension(target=Name(id='a', ctx=Store()), iter=Name(id='ab', ctx=Load()), ifs=[])]))"},
63+
{"( a for a, in ab )", "eval", "Expression(body=GeneratorExp(elt=Name(id='a', ctx=Load()), generators=[comprehension(target=Tuple(elts=[Name(id='a', ctx=Store())], ctx=Store()), iter=Name(id='ab', ctx=Load()), ifs=[])]))"},
64+
{"( a for a, b in ab )", "eval", "Expression(body=GeneratorExp(elt=Name(id='a', ctx=Load()), generators=[comprehension(target=Tuple(elts=[Name(id='a', ctx=Store()), Name(id='b', ctx=Store())], ctx=Store()), iter=Name(id='ab', ctx=Load()), ifs=[])]))"},
65+
{"( a for a in ab if a )", "eval", "Expression(body=GeneratorExp(elt=Name(id='a', ctx=Load()), generators=[comprehension(target=Name(id='a', ctx=Store()), iter=Name(id='ab', ctx=Load()), ifs=[Name(id='a', ctx=Load())])]))"},
66+
{"( a for a in ab if a if b if c )", "eval", "Expression(body=GeneratorExp(elt=Name(id='a', ctx=Load()), generators=[comprehension(target=Name(id='a', ctx=Store()), iter=Name(id='ab', ctx=Load()), ifs=[Name(id='a', ctx=Load()), Name(id='b', ctx=Load()), Name(id='c', ctx=Load())])]))"},
67+
{"( a for a in ab for A in AB )", "eval", "Expression(body=GeneratorExp(elt=Name(id='a', ctx=Load()), generators=[comprehension(target=Name(id='a', ctx=Store()), iter=Name(id='ab', ctx=Load()), ifs=[]), comprehension(target=Name(id='A', ctx=Store()), iter=Name(id='AB', ctx=Load()), ifs=[])]))"},
68+
{"( a for a in ab if a if b for A in AB if c )", "eval", "Expression(body=GeneratorExp(elt=Name(id='a', ctx=Load()), generators=[comprehension(target=Name(id='a', ctx=Store()), iter=Name(id='ab', ctx=Load()), ifs=[Name(id='a', ctx=Load()), Name(id='b', ctx=Load())]), comprehension(target=Name(id='A', ctx=Store()), iter=Name(id='AB', ctx=Load()), ifs=[Name(id='c', ctx=Load())])]))"},
69+
{"[ a for a in ab ]", "eval", "Expression(body=ListComp(elt=Name(id='a', ctx=Load()), generators=[comprehension(target=Name(id='a', ctx=Store()), iter=Name(id='ab', ctx=Load()), ifs=[])]))"},
70+
{"[ a for a, in ab ]", "eval", "Expression(body=ListComp(elt=Name(id='a', ctx=Load()), generators=[comprehension(target=Tuple(elts=[Name(id='a', ctx=Store())], ctx=Store()), iter=Name(id='ab', ctx=Load()), ifs=[])]))"},
71+
{"[ a for a, b in ab ]", "eval", "Expression(body=ListComp(elt=Name(id='a', ctx=Load()), generators=[comprehension(target=Tuple(elts=[Name(id='a', ctx=Store()), Name(id='b', ctx=Store())], ctx=Store()), iter=Name(id='ab', ctx=Load()), ifs=[])]))"},
72+
{"[ a for a in ab if a ]", "eval", "Expression(body=ListComp(elt=Name(id='a', ctx=Load()), generators=[comprehension(target=Name(id='a', ctx=Store()), iter=Name(id='ab', ctx=Load()), ifs=[Name(id='a', ctx=Load())])]))"},
73+
{"[ a for a in ab if a if b if c ]", "eval", "Expression(body=ListComp(elt=Name(id='a', ctx=Load()), generators=[comprehension(target=Name(id='a', ctx=Store()), iter=Name(id='ab', ctx=Load()), ifs=[Name(id='a', ctx=Load()), Name(id='b', ctx=Load()), Name(id='c', ctx=Load())])]))"},
74+
{"[ a for a in ab for A in AB ]", "eval", "Expression(body=ListComp(elt=Name(id='a', ctx=Load()), generators=[comprehension(target=Name(id='a', ctx=Store()), iter=Name(id='ab', ctx=Load()), ifs=[]), comprehension(target=Name(id='A', ctx=Store()), iter=Name(id='AB', ctx=Load()), ifs=[])]))"},
75+
{"[ a for a in ab if a if b for A in AB if c ]", "eval", "Expression(body=ListComp(elt=Name(id='a', ctx=Load()), generators=[comprehension(target=Name(id='a', ctx=Store()), iter=Name(id='ab', ctx=Load()), ifs=[Name(id='a', ctx=Load()), Name(id='b', ctx=Load())]), comprehension(target=Name(id='A', ctx=Store()), iter=Name(id='AB', ctx=Load()), ifs=[Name(id='c', ctx=Load())])]))"},
76+
{"{ a for a in ab }", "eval", "Expression(body=SetComp(elt=Name(id='a', ctx=Load()), generators=[comprehension(target=Name(id='a', ctx=Store()), iter=Name(id='ab', ctx=Load()), ifs=[])]))"},
77+
{"{ a for a, in ab }", "eval", "Expression(body=SetComp(elt=Name(id='a', ctx=Load()), generators=[comprehension(target=Tuple(elts=[Name(id='a', ctx=Store())], ctx=Store()), iter=Name(id='ab', ctx=Load()), ifs=[])]))"},
78+
{"{ a for a, b in ab }", "eval", "Expression(body=SetComp(elt=Name(id='a', ctx=Load()), generators=[comprehension(target=Tuple(elts=[Name(id='a', ctx=Store()), Name(id='b', ctx=Store())], ctx=Store()), iter=Name(id='ab', ctx=Load()), ifs=[])]))"},
79+
{"{ a for a in ab if a }", "eval", "Expression(body=SetComp(elt=Name(id='a', ctx=Load()), generators=[comprehension(target=Name(id='a', ctx=Store()), iter=Name(id='ab', ctx=Load()), ifs=[Name(id='a', ctx=Load())])]))"},
80+
{"{ a for a in ab if a if b if c }", "eval", "Expression(body=SetComp(elt=Name(id='a', ctx=Load()), generators=[comprehension(target=Name(id='a', ctx=Store()), iter=Name(id='ab', ctx=Load()), ifs=[Name(id='a', ctx=Load()), Name(id='b', ctx=Load()), Name(id='c', ctx=Load())])]))"},
81+
{"{ a for a in ab for A in AB }", "eval", "Expression(body=SetComp(elt=Name(id='a', ctx=Load()), generators=[comprehension(target=Name(id='a', ctx=Store()), iter=Name(id='ab', ctx=Load()), ifs=[]), comprehension(target=Name(id='A', ctx=Store()), iter=Name(id='AB', ctx=Load()), ifs=[])]))"},
82+
{"{ a for a in ab if a if b for A in AB if c }", "eval", "Expression(body=SetComp(elt=Name(id='a', ctx=Load()), generators=[comprehension(target=Name(id='a', ctx=Store()), iter=Name(id='ab', ctx=Load()), ifs=[Name(id='a', ctx=Load()), Name(id='b', ctx=Load())]), comprehension(target=Name(id='A', ctx=Store()), iter=Name(id='AB', ctx=Load()), ifs=[Name(id='c', ctx=Load())])]))"},
83+
{"{ a:b for a in ab }", "eval", "Expression(body=DictComp(key=Name(id='a', ctx=Load()), value=Name(id='b', ctx=Load()), generators=[comprehension(target=Name(id='a', ctx=Store()), iter=Name(id='ab', ctx=Load()), ifs=[])]))"},
84+
{"{ a:b for a, in ab }", "eval", "Expression(body=DictComp(key=Name(id='a', ctx=Load()), value=Name(id='b', ctx=Load()), generators=[comprehension(target=Tuple(elts=[Name(id='a', ctx=Store())], ctx=Store()), iter=Name(id='ab', ctx=Load()), ifs=[])]))"},
85+
{"{ a:b for a, b in ab }", "eval", "Expression(body=DictComp(key=Name(id='a', ctx=Load()), value=Name(id='b', ctx=Load()), generators=[comprehension(target=Tuple(elts=[Name(id='a', ctx=Store()), Name(id='b', ctx=Store())], ctx=Store()), iter=Name(id='ab', ctx=Load()), ifs=[])]))"},
86+
{"{ a:b for a in ab if a }", "eval", "Expression(body=DictComp(key=Name(id='a', ctx=Load()), value=Name(id='b', ctx=Load()), generators=[comprehension(target=Name(id='a', ctx=Store()), iter=Name(id='ab', ctx=Load()), ifs=[Name(id='a', ctx=Load())])]))"},
87+
{"{ a:b for a in ab if a if b if c }", "eval", "Expression(body=DictComp(key=Name(id='a', ctx=Load()), value=Name(id='b', ctx=Load()), generators=[comprehension(target=Name(id='a', ctx=Store()), iter=Name(id='ab', ctx=Load()), ifs=[Name(id='a', ctx=Load()), Name(id='b', ctx=Load()), Name(id='c', ctx=Load())])]))"},
88+
{"{ a:b for a in ab for A in AB }", "eval", "Expression(body=DictComp(key=Name(id='a', ctx=Load()), value=Name(id='b', ctx=Load()), generators=[comprehension(target=Name(id='a', ctx=Store()), iter=Name(id='ab', ctx=Load()), ifs=[]), comprehension(target=Name(id='A', ctx=Store()), iter=Name(id='AB', ctx=Load()), ifs=[])]))"},
89+
{"{ a:b for a in ab if a if b for A in AB if c }", "eval", "Expression(body=DictComp(key=Name(id='a', ctx=Load()), value=Name(id='b', ctx=Load()), generators=[comprehension(target=Name(id='a', ctx=Store()), iter=Name(id='ab', ctx=Load()), ifs=[Name(id='a', ctx=Load()), Name(id='b', ctx=Load())]), comprehension(target=Name(id='A', ctx=Store()), iter=Name(id='AB', ctx=Load()), ifs=[Name(id='c', ctx=Load())])]))"},
6290
// END TESTS
6391
} {
6492
Ast, err := ParseString(test.in, test.mode)
@@ -67,7 +95,7 @@ func TestGrammar(t *testing.T) {
6795
} else {
6896
out := ast.Dump(Ast)
6997
if out != test.out {
70-
t.Errorf("Parse(%q) expecting %q actual %q", test.in, test.out, out)
98+
t.Errorf("Parse(%q)\nwant> %q\n got> %q\n", test.in, test.out, out)
7199
}
72100
}
73101
}

‎parser/make_grammar_test.py

+33
Original file line numberDiff line numberDiff line change
@@ -46,6 +46,39 @@
4646
("[1,]", "eval"),
4747
("[1,2]", "eval"),
4848
("[1,2,]", "eval"),
49+
("( a for a in ab )", "eval"),
50+
("( a for a, in ab )", "eval"),
51+
("( a for a, b in ab )", "eval"),
52+
("( a for a in ab if a )", "eval"),
53+
("( a for a in ab if a if b if c )", "eval"),
54+
("( a for a in ab for A in AB )", "eval"),
55+
("( a for a in ab if a if b for A in AB if c )", "eval"),
56+
57+
("[ a for a in ab ]", "eval"),
58+
("[ a for a, in ab ]", "eval"),
59+
("[ a for a, b in ab ]", "eval"),
60+
("[ a for a in ab if a ]", "eval"),
61+
("[ a for a in ab if a if b if c ]", "eval"),
62+
("[ a for a in ab for A in AB ]", "eval"),
63+
("[ a for a in ab if a if b for A in AB if c ]", "eval"),
64+
65+
("{ a for a in ab }", "eval"),
66+
("{ a for a, in ab }", "eval"),
67+
("{ a for a, b in ab }", "eval"),
68+
("{ a for a in ab if a }", "eval"),
69+
("{ a for a in ab if a if b if c }", "eval"),
70+
("{ a for a in ab for A in AB }", "eval"),
71+
("{ a for a in ab if a if b for A in AB if c }", "eval"),
72+
73+
("{ a:b for a in ab }", "eval"),
74+
("{ a:b for a, in ab }", "eval"),
75+
("{ a:b for a, b in ab }", "eval"),
76+
("{ a:b for a in ab if a }", "eval"),
77+
("{ a:b for a in ab if a if b if c }", "eval"),
78+
("{ a:b for a in ab for A in AB }", "eval"),
79+
("{ a:b for a in ab if a if b for A in AB if c }", "eval"),
80+
81+
# ("del a,b", "exec"),
4982
]
5083

5184
def dump(source, mode):

0 commit comments

Comments
 (0)