Skip to content

Commit 952d70b

Browse files
charliermarshcharliermarsh
committed
Add expression context parsing
1 parent 02953b9 commit 952d70b

29 files changed

Lines changed: 1024 additions & 581 deletions

File tree

‎parser/python.lalrpop‎

Lines changed: 9 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,7 @@ use crate::{
88
error::{LexicalError, LexicalErrorType},
99
function::{ArgumentList, parse_args, parse_params},
1010
lexer,
11+
context::set_context,
1112
string::parse_strings,
1213
token::StringKind
1314
};
@@ -82,7 +83,7 @@ DelStatement: ast::Stmt = {
8283
location,
8384
end_location: Some(end_location),
8485
custom: (),
85-
node: ast::StmtKind::Delete { targets },
86+
node: ast::StmtKind::Delete { targets: targets.into_iter().map(|expr| set_context(expr, ast::ExprContext::Del)).collect() },
8687
}
8788
},
8889
};
@@ -98,11 +99,11 @@ ExpressionStatement: ast::Stmt = {
9899
node: ast::StmtKind::Expr { value: Box::new(expression) }
99100
}
100101
} else {
101-
let mut targets = vec![expression];
102+
let mut targets = vec![set_context(expression, ast::ExprContext::Store)];
102103
let mut values = suffix;
103104

104105
while values.len() > 1 {
105-
targets.push(values.remove(0));
106+
targets.push(set_context(values.remove(0), ast::ExprContext::Store));
106107
}
107108

108109
let value = Box::new(values.into_iter().next().unwrap());
@@ -121,7 +122,7 @@ ExpressionStatement: ast::Stmt = {
121122
location,
122123
end_location: Some(end_location),
123124
node: ast::StmtKind::AugAssign {
124-
target: Box::new(target),
125+
target: Box::new(set_context(target, ast::ExprContext::Store)),
125126
op,
126127
value: Box::new(rhs)
127128
},
@@ -134,7 +135,7 @@ ExpressionStatement: ast::Stmt = {
134135
location,
135136
end_location: Some(end_location),
136137
node: ast::StmtKind::AnnAssign {
137-
target: Box::new(target),
138+
target: Box::new(set_context(target, ast::ExprContext::Store)),
138139
annotation: Box::new(annotation),
139140
value: rhs.map(Box::new),
140141
simple: if simple { 1 } else { 0 },
@@ -399,7 +400,7 @@ WhileStatement: ast::Stmt = {
399400
ForStatement: ast::Stmt = {
400401
<location:@L> <is_async:"async"?> "for" <target:ExpressionList> "in" <iter:TestList> ":" <body:Suite> <s2:("else" ":" Suite)?> <end_location:@R> => {
401402
let orelse = s2.map(|s| s.2).unwrap_or_default();
402-
let target = Box::new(target);
403+
let target = Box::new(set_context(target, ast::ExprContext::Store));
403404
let iter = Box::new(iter);
404405
let type_comment = None;
405406
let node = if is_async.is_some() {
@@ -484,7 +485,7 @@ WithStatement: ast::Stmt = {
484485

485486
WithItem: ast::Withitem = {
486487
<context_expr:Test> <n:("as" Expression)?> => {
487-
let optional_vars = n.map(|val| Box::new(val.1));
488+
let optional_vars = n.map(|val| Box::new(set_context(val.1, ast::ExprContext::Store)));
488489
let context_expr = Box::new(context_expr);
489490
ast::Withitem { context_expr, optional_vars }
490491
},
@@ -1233,7 +1234,7 @@ SingleForComprehension: ast::Comprehension = {
12331234
<location:@L> <is_async:"async"?> "for" <target:ExpressionList> "in" <iter:OrTest> <ifs:ComprehensionIf*> <end_location:@R> => {
12341235
let is_async = is_async.is_some();
12351236
ast::Comprehension {
1236-
target: Box::new(target),
1237+
target: Box::new(set_context(target, ast::ExprContext::Store)),
12371238
iter: Box::new(iter),
12381239
ifs,
12391240
is_async: if is_async { 1 } else { 0 },

‎parser/src/context.rs‎

Lines changed: 177 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,177 @@
1+
use rustpython_ast::{Expr, ExprContext, ExprKind};
2+
3+
pub fn set_context(expr: Expr, ctx: ExprContext) -> Expr {
4+
match expr.node {
5+
ExprKind::Name { id, .. } => Expr {
6+
node: ExprKind::Name { id, ctx },
7+
..expr
8+
},
9+
ExprKind::Tuple { elts, .. } => Expr {
10+
node: ExprKind::Tuple {
11+
elts: elts
12+
.into_iter()
13+
.map(|elt| set_context(elt, ctx.clone()))
14+
.collect(),
15+
ctx,
16+
},
17+
..expr
18+
},
19+
ExprKind::List { elts, .. } => Expr {
20+
node: ExprKind::List {
21+
elts: elts
22+
.into_iter()
23+
.map(|elt| set_context(elt, ctx.clone()))
24+
.collect(),
25+
ctx,
26+
},
27+
..expr
28+
},
29+
ExprKind::Attribute { value, attr, .. } => Expr {
30+
node: ExprKind::Attribute { value, attr, ctx },
31+
..expr
32+
},
33+
ExprKind::Subscript { value, slice, .. } => Expr {
34+
node: ExprKind::Subscript { value, slice, ctx },
35+
..expr
36+
},
37+
ExprKind::Starred { value, .. } => Expr {
38+
node: ExprKind::Starred {
39+
value: Box::new(set_context(*value, ctx.clone())),
40+
ctx,
41+
},
42+
..expr
43+
},
44+
_ => expr,
45+
}
46+
}
47+
48+
#[cfg(test)]
49+
mod tests {
50+
use crate::parser::parse_program;
51+
52+
#[test]
53+
fn test_assign_name() {
54+
let source = String::from("x = (1, 2, 3)");
55+
let parse_ast = parse_program(&source, "<test>").unwrap();
56+
insta::assert_debug_snapshot!(parse_ast);
57+
}
58+
59+
#[test]
60+
fn test_assign_tuple() {
61+
let source = String::from("(x, y) = (1, 2, 3)");
62+
let parse_ast = parse_program(&source, "<test>").unwrap();
63+
insta::assert_debug_snapshot!(parse_ast);
64+
}
65+
66+
#[test]
67+
fn test_assign_list() {
68+
let source = String::from("[x, y] = (1, 2, 3)");
69+
let parse_ast = parse_program(&source, "<test>").unwrap();
70+
insta::assert_debug_snapshot!(parse_ast);
71+
}
72+
73+
#[test]
74+
fn test_assign_attribute() {
75+
let source = String::from("x.y = (1, 2, 3)");
76+
let parse_ast = parse_program(&source, "<test>").unwrap();
77+
insta::assert_debug_snapshot!(parse_ast);
78+
}
79+
80+
#[test]
81+
fn test_assign_subscript() {
82+
let source = String::from("x[y] = (1, 2, 3)");
83+
let parse_ast = parse_program(&source, "<test>").unwrap();
84+
insta::assert_debug_snapshot!(parse_ast);
85+
}
86+
87+
#[test]
88+
fn test_assign_starred() {
89+
let source = String::from("(x, *y) = (1, 2, 3)");
90+
let parse_ast = parse_program(&source, "<test>").unwrap();
91+
insta::assert_debug_snapshot!(parse_ast);
92+
}
93+
94+
#[test]
95+
fn test_assign_for() {
96+
let source = String::from("for x in (1, 2, 3): pass");
97+
let parse_ast = parse_program(&source, "<test>").unwrap();
98+
insta::assert_debug_snapshot!(parse_ast);
99+
}
100+
101+
#[test]
102+
fn test_assign_list_comp() {
103+
let source = String::from("x = [y for y in (1, 2, 3)]");
104+
let parse_ast = parse_program(&source, "<test>").unwrap();
105+
insta::assert_debug_snapshot!(parse_ast);
106+
}
107+
108+
#[test]
109+
fn test_assign_set_comp() {
110+
let source = String::from("x = {y for y in (1, 2, 3)}");
111+
let parse_ast = parse_program(&source, "<test>").unwrap();
112+
insta::assert_debug_snapshot!(parse_ast);
113+
}
114+
115+
#[test]
116+
fn test_assign_with() {
117+
let source = String::from("with 1 as x: pass");
118+
let parse_ast = parse_program(&source, "<test>").unwrap();
119+
insta::assert_debug_snapshot!(parse_ast);
120+
}
121+
122+
#[test]
123+
fn test_assign_named_expr() {
124+
let source = String::from("if x:= 1: pass");
125+
let parse_ast = parse_program(&source, "<test>").unwrap();
126+
insta::assert_debug_snapshot!(parse_ast);
127+
}
128+
129+
#[test]
130+
fn test_ann_assign_name() {
131+
let source = String::from("x: int = 1");
132+
let parse_ast = parse_program(&source, "<test>").unwrap();
133+
insta::assert_debug_snapshot!(parse_ast);
134+
}
135+
136+
#[test]
137+
fn test_aug_assign_name() {
138+
let source = String::from("x += 1");
139+
let parse_ast = parse_program(&source, "<test>").unwrap();
140+
insta::assert_debug_snapshot!(parse_ast);
141+
}
142+
143+
#[test]
144+
fn test_aug_assign_attribute() {
145+
let source = String::from("x.y += (1, 2, 3)");
146+
let parse_ast = parse_program(&source, "<test>").unwrap();
147+
insta::assert_debug_snapshot!(parse_ast);
148+
}
149+
150+
#[test]
151+
fn test_aug_assign_subscript() {
152+
let source = String::from("x[y] += (1, 2, 3)");
153+
let parse_ast = parse_program(&source, "<test>").unwrap();
154+
insta::assert_debug_snapshot!(parse_ast);
155+
}
156+
157+
#[test]
158+
fn test_del_name() {
159+
let source = String::from("del x");
160+
let parse_ast = parse_program(&source, "<test>").unwrap();
161+
insta::assert_debug_snapshot!(parse_ast);
162+
}
163+
164+
#[test]
165+
fn test_del_attribute() {
166+
let source = String::from("del x.y");
167+
let parse_ast = parse_program(&source, "<test>").unwrap();
168+
insta::assert_debug_snapshot!(parse_ast);
169+
}
170+
171+
#[test]
172+
fn test_del_subscript() {
173+
let source = String::from("del x[y]");
174+
let parse_ast = parse_program(&source, "<test>").unwrap();
175+
insta::assert_debug_snapshot!(parse_ast);
176+
}
177+
}

‎parser/src/lib.rs‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -30,5 +30,6 @@ pub mod mode;
3030
pub mod parser;
3131
#[rustfmt::skip]
3232
mod python;
33+
mod context;
3334
mod string;
3435
pub mod token;

‎parser/src/snapshots/rustpython_parser__context__tests__ann_assign_name.snap‎

Lines changed: 28 additions & 20 deletions
Original file line numberDiff line numberDiff line change
@@ -4,40 +4,46 @@ expression: parse_ast
44
---
55
[
66
Located {
7-
start: Location {
7+
location: Location {
88
row: 1,
99
column: 1,
1010
},
11-
end: Location {
12-
row: 1,
13-
column: 11,
14-
},
11+
end_location: Some(
12+
Location {
13+
row: 1,
14+
column: 11,
15+
},
16+
),
1517
custom: (),
1618
node: AnnAssign {
1719
target: Located {
18-
start: Location {
20+
location: Location {
1921
row: 1,
2022
column: 1,
2123
},
22-
end: Location {
23-
row: 1,
24-
column: 2,
25-
},
24+
end_location: Some(
25+
Location {
26+
row: 1,
27+
column: 2,
28+
},
29+
),
2630
custom: (),
2731
node: Name {
2832
id: "x",
2933
ctx: Store,
3034
},
3135
},
3236
annotation: Located {
33-
start: Location {
37+
location: Location {
3438
row: 1,
3539
column: 4,
3640
},
37-
end: Location {
38-
row: 1,
39-
column: 7,
40-
},
41+
end_location: Some(
42+
Location {
43+
row: 1,
44+
column: 7,
45+
},
46+
),
4147
custom: (),
4248
node: Name {
4349
id: "int",
@@ -46,14 +52,16 @@ expression: parse_ast
4652
},
4753
value: Some(
4854
Located {
49-
start: Location {
55+
location: Location {
5056
row: 1,
5157
column: 10,
5258
},
53-
end: Location {
54-
row: 1,
55-
column: 11,
56-
},
59+
end_location: Some(
60+
Location {
61+
row: 1,
62+
column: 11,
63+
},
64+
),
5765
custom: (),
5866
node: Constant {
5967
value: Int(

‎parser/src/snapshots/rustpython_parser__context__tests__assign.snap‎

Lines changed: 42 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,42 @@
1+
---
2+
source: compiler/parser/src/context.rs
3+
expression: parse_ast
4+
---
5+
[
6+
Located {
7+
location: Location {
8+
row: 1,
9+
column: 1,
10+
},
11+
custom: (),
12+
node: Assign {
13+
targets: [
14+
Located {
15+
location: Location {
16+
row: 1,
17+
column: 1,
18+
},
19+
custom: (),
20+
node: Name {
21+
id: "x",
22+
ctx: Store,
23+
},
24+
},
25+
],
26+
value: Located {
27+
location: Location {
28+
row: 1,
29+
column: 5,
30+
},
31+
custom: (),
32+
node: Constant {
33+
value: Int(
34+
1,
35+
),
36+
kind: None,
37+
},
38+
},
39+
type_comment: None,
40+
},
41+
},
42+
]

0 commit comments

Comments
 (0)