Skip to content

Commit bad0904

Browse files
committed
Assign and augmented assign
1 parent 8858589 commit bad0904

File tree

3 files changed

+59
-6
lines changed

3 files changed

+59
-6
lines changed

parser/grammar.y

Lines changed: 21 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -42,6 +42,13 @@ func applyTrailers(expr ast.Expr, trailers []ast.Expr) ast.Expr {
4242
return expr
4343
}
4444

45+
// Set the context for all the items in exprs
46+
func setCtx(exprs []ast.Expr, ctx ast.ExprContext) {
47+
for i := range exprs {
48+
exprs[i].(ast.SetCtxer).SetCtx(ctx)
49+
}
50+
}
51+
4552
%}
4653

4754
%union {
@@ -77,7 +84,7 @@ func applyTrailers(expr ast.Expr, trailers []ast.Expr) ast.Expr {
7784
%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 import_name import_from while_stmt if_stmt for_stmt try_stmt with_stmt
7885
%type <op> augassign
7986
%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 sliceop arglist except_clause
80-
%type <exprs> exprlist testlistraw comp_if comp_iter expr_or_star_exprs test_or_star_exprs tests test_colon_tests trailers
87+
%type <exprs> exprlist testlistraw comp_if comp_iter expr_or_star_exprs test_or_star_exprs tests test_colon_tests trailers equals_yield_expr_or_testlist_star_expr
8188
%type <cmpop> comp_op
8289
%type <comma> optional_comma
8390
%type <comprehensions> comp_for
@@ -553,11 +560,18 @@ expr_stmt: testlist_star_expr ('=' (yield_expr|testlist_star_expr))*
553560
expr_stmt:
554561
testlist_star_expr augassign yield_expr_or_testlist
555562
{
556-
// FIXME
563+
target := $1
564+
target.(ast.SetCtxer).SetCtx(ast.Store)
565+
$$ = &ast.AugAssign{StmtBase: ast.StmtBase{$<pos>$}, Target: target, Op: $2, Value: $3}
557566
}
558567
| testlist_star_expr equals_yield_expr_or_testlist_star_expr
559568
{
560-
// FIXME
569+
targets := []ast.Expr{$1}
570+
targets = append(targets, $2...)
571+
value := targets[len(targets)-1]
572+
targets = targets[:len(targets)-1]
573+
setCtx(targets, ast.Store)
574+
$$ = &ast.Assign{StmtBase: ast.StmtBase{$<pos>$}, Targets: targets, Value: value}
561575
}
562576
| testlist_star_expr
563577
{
@@ -587,9 +601,12 @@ yield_expr_or_testlist_star_expr:
587601
equals_yield_expr_or_testlist_star_expr:
588602
'=' yield_expr_or_testlist_star_expr
589603
{
604+
$$ = nil
605+
$$ = append($$, $2)
590606
}
591607
| equals_yield_expr_or_testlist_star_expr '=' yield_expr_or_testlist_star_expr
592608
{
609+
$$ = append($$, $3)
593610
}
594611

595612
test_or_star_exprs:
@@ -682,9 +699,7 @@ augassign:
682699
del_stmt:
683700
DEL exprlist
684701
{
685-
for i := range $2 {
686-
$2[i].(ast.SetCtxer).SetCtx(ast.Del)
687-
}
702+
setCtx($2, ast.Del)
688703
$$ = &ast.Delete{StmtBase: ast.StmtBase{$<pos>$}, Targets: $2}
689704
}
690705

parser/grammar_test.go

Lines changed: 17 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -217,6 +217,23 @@ func TestGrammar(t *testing.T) {
217217
{"with x:\n pass\n", "exec", "Module(body=[With(items=[withitem(context_expr=Name(id='x', ctx=Load()), optional_vars=None)], body=[Pass()])])"},
218218
{"with x as y:\n pass\n", "exec", "Module(body=[With(items=[withitem(context_expr=Name(id='x', ctx=Load()), optional_vars=Name(id='y', ctx=Store()))], body=[Pass()])])"},
219219
{"with x as y, a as b, c, d as e:\n pass\n continue\n", "exec", "Module(body=[With(items=[withitem(context_expr=Name(id='x', ctx=Load()), optional_vars=Name(id='y', ctx=Store())), withitem(context_expr=Name(id='a', ctx=Load()), optional_vars=Name(id='b', ctx=Store())), withitem(context_expr=Name(id='c', ctx=Load()), optional_vars=None), withitem(context_expr=Name(id='d', ctx=Load()), optional_vars=Name(id='e', ctx=Store()))], body=[Pass(), Continue()])])"},
220+
{"a += b", "exec", "Module(body=[AugAssign(target=Name(id='a', ctx=Store()), op=Add(), value=Name(id='b', ctx=Load()))])"},
221+
{"a -= b", "exec", "Module(body=[AugAssign(target=Name(id='a', ctx=Store()), op=Sub(), value=Name(id='b', ctx=Load()))])"},
222+
{"a *= b", "exec", "Module(body=[AugAssign(target=Name(id='a', ctx=Store()), op=Mult(), value=Name(id='b', ctx=Load()))])"},
223+
{"a /= b", "exec", "Module(body=[AugAssign(target=Name(id='a', ctx=Store()), op=Div(), value=Name(id='b', ctx=Load()))])"},
224+
{"a -= b", "exec", "Module(body=[AugAssign(target=Name(id='a', ctx=Store()), op=Sub(), value=Name(id='b', ctx=Load()))])"},
225+
{"a %= b", "exec", "Module(body=[AugAssign(target=Name(id='a', ctx=Store()), op=Mod(), value=Name(id='b', ctx=Load()))])"},
226+
{"a &= b", "exec", "Module(body=[AugAssign(target=Name(id='a', ctx=Store()), op=BitAnd(), value=Name(id='b', ctx=Load()))])"},
227+
{"a |= b", "exec", "Module(body=[AugAssign(target=Name(id='a', ctx=Store()), op=BitOr(), value=Name(id='b', ctx=Load()))])"},
228+
{"a ^= b", "exec", "Module(body=[AugAssign(target=Name(id='a', ctx=Store()), op=BitXor(), value=Name(id='b', ctx=Load()))])"},
229+
{"a <<= b", "exec", "Module(body=[AugAssign(target=Name(id='a', ctx=Store()), op=LShift(), value=Name(id='b', ctx=Load()))])"},
230+
{"a >>= b", "exec", "Module(body=[AugAssign(target=Name(id='a', ctx=Store()), op=RShift(), value=Name(id='b', ctx=Load()))])"},
231+
{"a **= b", "exec", "Module(body=[AugAssign(target=Name(id='a', ctx=Store()), op=Pow(), value=Name(id='b', ctx=Load()))])"},
232+
{"a //= b", "exec", "Module(body=[AugAssign(target=Name(id='a', ctx=Store()), op=FloorDiv(), value=Name(id='b', ctx=Load()))])"},
233+
{"a = b", "exec", "Module(body=[Assign(targets=[Name(id='a', ctx=Store())], value=Name(id='b', ctx=Load()))])"},
234+
{"a = b = c", "exec", "Module(body=[Assign(targets=[Name(id='a', ctx=Store()), Name(id='b', ctx=Store())], value=Name(id='c', ctx=Load()))])"},
235+
{"a, b = 1, 2", "exec", "Module(body=[Assign(targets=[Tuple(elts=[Name(id='a', ctx=Store()), Name(id='b', ctx=Store())], ctx=Store())], value=Tuple(elts=[Num(n=1), Num(n=2)], ctx=Load()))])"},
236+
{"a, b = c, d = 1, 2", "exec", "Module(body=[Assign(targets=[Tuple(elts=[Name(id='a', ctx=Store()), Name(id='b', ctx=Store())], ctx=Store()), Tuple(elts=[Name(id='c', ctx=Store()), Name(id='d', ctx=Store())], ctx=Store())], value=Tuple(elts=[Num(n=1), Num(n=2)], ctx=Load()))])"},
220237
// END TESTS
221238
} {
222239
Ast, err := ParseString(test.in, test.mode)

parser/make_grammar_test.py

Lines changed: 21 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -314,6 +314,27 @@
314314
pass
315315
continue
316316
""", "exec"),
317+
318+
# Augmented assign
319+
("a += b", "exec"),
320+
("a -= b", "exec"),
321+
("a *= b", "exec"),
322+
("a /= b", "exec"),
323+
("a -= b", "exec"),
324+
("a %= b", "exec"),
325+
("a &= b", "exec"),
326+
("a |= b", "exec"),
327+
("a ^= b", "exec"),
328+
("a <<= b", "exec"),
329+
("a >>= b", "exec"),
330+
("a **= b", "exec"),
331+
("a //= b", "exec"),
332+
333+
# Assign
334+
("a = b", "exec"),
335+
("a = b = c", "exec"),
336+
("a, b = 1, 2", "exec"),
337+
("a, b = c, d = 1, 2", "exec"),
317338
]
318339

319340
def dump(source, mode):

0 commit comments

Comments
 (0)