Add expression context parsing

astral-sh · Oct 17, 2022 · 952d70b · 952d70b
1 parent 02953b9
commit 952d70b
Show file tree

Hide file tree

Showing 29 changed files with 1,024 additions and 581 deletions.
diff --git a/parser/python.lalrpop b/parser/python.lalrpop
@@ -8,6 +8,7 @@ use crate::{
     error::{LexicalError, LexicalErrorType},
     function::{ArgumentList, parse_args, parse_params},
     lexer,
+    context::set_context,
     string::parse_strings,
     token::StringKind
 };
@@ -82,7 +83,7 @@ DelStatement: ast::Stmt = {
             location,
             end_location: Some(end_location),
             custom: (),
-            node: ast::StmtKind::Delete { targets },
+            node: ast::StmtKind::Delete { targets: targets.into_iter().map(|expr| set_context(expr, ast::ExprContext::Del)).collect() },
         }
     },
 };
@@ -98,11 +99,11 @@ ExpressionStatement: ast::Stmt = {
                 node: ast::StmtKind::Expr { value: Box::new(expression) }
             }
         } else {
-            let mut targets = vec![expression];
+            let mut targets = vec![set_context(expression, ast::ExprContext::Store)];
             let mut values = suffix;
 
             while values.len() > 1 {
-                targets.push(values.remove(0));
+                targets.push(set_context(values.remove(0), ast::ExprContext::Store));
             }
 
             let value = Box::new(values.into_iter().next().unwrap());
@@ -121,7 +122,7 @@ ExpressionStatement: ast::Stmt = {
             location,
             end_location: Some(end_location),
             node: ast::StmtKind::AugAssign {
-                target: Box::new(target),
+                target: Box::new(set_context(target, ast::ExprContext::Store)),
                 op,
                 value: Box::new(rhs)
             },
@@ -134,7 +135,7 @@ ExpressionStatement: ast::Stmt = {
             location,
             end_location: Some(end_location),
             node: ast::StmtKind::AnnAssign {
-                target: Box::new(target),
+                target: Box::new(set_context(target, ast::ExprContext::Store)),
                 annotation: Box::new(annotation),
                 value: rhs.map(Box::new),
                 simple: if simple { 1 } else { 0 },
@@ -399,7 +400,7 @@ WhileStatement: ast::Stmt = {
 ForStatement: ast::Stmt = {
     <location:@L> <is_async:"async"?> "for" <target:ExpressionList> "in" <iter:TestList> ":" <body:Suite> <s2:("else" ":" Suite)?> <end_location:@R> => {
         let orelse = s2.map(|s| s.2).unwrap_or_default();
-        let target = Box::new(target);
+        let target = Box::new(set_context(target, ast::ExprContext::Store));
         let iter = Box::new(iter);
         let type_comment = None;
         let node = if is_async.is_some() {
@@ -484,7 +485,7 @@ WithStatement: ast::Stmt = {
 
 WithItem: ast::Withitem = {
     <context_expr:Test> <n:("as" Expression)?> => {
-        let optional_vars = n.map(|val| Box::new(val.1));
+        let optional_vars = n.map(|val| Box::new(set_context(val.1, ast::ExprContext::Store)));
         let context_expr = Box::new(context_expr);
         ast::Withitem { context_expr, optional_vars }
     },
@@ -1233,7 +1234,7 @@ SingleForComprehension: ast::Comprehension = {
     <location:@L> <is_async:"async"?> "for" <target:ExpressionList> "in" <iter:OrTest> <ifs:ComprehensionIf*> <end_location:@R> => {
         let is_async = is_async.is_some();
         ast::Comprehension {
-            target: Box::new(target),
+            target: Box::new(set_context(target, ast::ExprContext::Store)),
             iter: Box::new(iter),
             ifs,
             is_async: if is_async { 1 } else { 0 },

diff --git a/parser/src/context.rs b/parser/src/context.rs
@@ -0,0 +1,177 @@
+use rustpython_ast::{Expr, ExprContext, ExprKind};
+
+pub fn set_context(expr: Expr, ctx: ExprContext) -> Expr {
+    match expr.node {
+        ExprKind::Name { id, .. } => Expr {
+            node: ExprKind::Name { id, ctx },
+            ..expr
+        },
+        ExprKind::Tuple { elts, .. } => Expr {
+            node: ExprKind::Tuple {
+                elts: elts
+                    .into_iter()
+                    .map(|elt| set_context(elt, ctx.clone()))
+                    .collect(),
+                ctx,
+            },
+            ..expr
+        },
+        ExprKind::List { elts, .. } => Expr {
+            node: ExprKind::List {
+                elts: elts
+                    .into_iter()
+                    .map(|elt| set_context(elt, ctx.clone()))
+                    .collect(),
+                ctx,
+            },
+            ..expr
+        },
+        ExprKind::Attribute { value, attr, .. } => Expr {
+            node: ExprKind::Attribute { value, attr, ctx },
+            ..expr
+        },
+        ExprKind::Subscript { value, slice, .. } => Expr {
+            node: ExprKind::Subscript { value, slice, ctx },
+            ..expr
+        },
+        ExprKind::Starred { value, .. } => Expr {
+            node: ExprKind::Starred {
+                value: Box::new(set_context(*value, ctx.clone())),
+                ctx,
+            },
+            ..expr
+        },
+        _ => expr,
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use crate::parser::parse_program;
+
+    #[test]
+    fn test_assign_name() {
+        let source = String::from("x = (1, 2, 3)");
+        let parse_ast = parse_program(&source, "<test>").unwrap();
+        insta::assert_debug_snapshot!(parse_ast);
+    }
+
+    #[test]
+    fn test_assign_tuple() {
+        let source = String::from("(x, y) = (1, 2, 3)");
+        let parse_ast = parse_program(&source, "<test>").unwrap();
+        insta::assert_debug_snapshot!(parse_ast);
+    }
+
+    #[test]
+    fn test_assign_list() {
+        let source = String::from("[x, y] = (1, 2, 3)");
+        let parse_ast = parse_program(&source, "<test>").unwrap();
+        insta::assert_debug_snapshot!(parse_ast);
+    }
+
+    #[test]
+    fn test_assign_attribute() {
+        let source = String::from("x.y = (1, 2, 3)");
+        let parse_ast = parse_program(&source, "<test>").unwrap();
+        insta::assert_debug_snapshot!(parse_ast);
+    }
+
+    #[test]
+    fn test_assign_subscript() {
+        let source = String::from("x[y] = (1, 2, 3)");
+        let parse_ast = parse_program(&source, "<test>").unwrap();
+        insta::assert_debug_snapshot!(parse_ast);
+    }
+
+    #[test]
+    fn test_assign_starred() {
+        let source = String::from("(x, *y) = (1, 2, 3)");
+        let parse_ast = parse_program(&source, "<test>").unwrap();
+        insta::assert_debug_snapshot!(parse_ast);
+    }
+
+    #[test]
+    fn test_assign_for() {
+        let source = String::from("for x in (1, 2, 3): pass");
+        let parse_ast = parse_program(&source, "<test>").unwrap();
+        insta::assert_debug_snapshot!(parse_ast);
+    }
+
+    #[test]
+    fn test_assign_list_comp() {
+        let source = String::from("x = [y for y in (1, 2, 3)]");
+        let parse_ast = parse_program(&source, "<test>").unwrap();
+        insta::assert_debug_snapshot!(parse_ast);
+    }
+
+    #[test]
+    fn test_assign_set_comp() {
+        let source = String::from("x = {y for y in (1, 2, 3)}");
+        let parse_ast = parse_program(&source, "<test>").unwrap();
+        insta::assert_debug_snapshot!(parse_ast);
+    }
+
+    #[test]
+    fn test_assign_with() {
+        let source = String::from("with 1 as x: pass");
+        let parse_ast = parse_program(&source, "<test>").unwrap();
+        insta::assert_debug_snapshot!(parse_ast);
+    }
+
+    #[test]
+    fn test_assign_named_expr() {
+        let source = String::from("if x:= 1: pass");
+        let parse_ast = parse_program(&source, "<test>").unwrap();
+        insta::assert_debug_snapshot!(parse_ast);
+    }
+
+    #[test]
+    fn test_ann_assign_name() {
+        let source = String::from("x: int = 1");
+        let parse_ast = parse_program(&source, "<test>").unwrap();
+        insta::assert_debug_snapshot!(parse_ast);
+    }
+
+    #[test]
+    fn test_aug_assign_name() {
+        let source = String::from("x += 1");
+        let parse_ast = parse_program(&source, "<test>").unwrap();
+        insta::assert_debug_snapshot!(parse_ast);
+    }
+
+    #[test]
+    fn test_aug_assign_attribute() {
+        let source = String::from("x.y += (1, 2, 3)");
+        let parse_ast = parse_program(&source, "<test>").unwrap();
+        insta::assert_debug_snapshot!(parse_ast);
+    }
+
+    #[test]
+    fn test_aug_assign_subscript() {
+        let source = String::from("x[y] += (1, 2, 3)");
+        let parse_ast = parse_program(&source, "<test>").unwrap();
+        insta::assert_debug_snapshot!(parse_ast);
+    }
+
+    #[test]
+    fn test_del_name() {
+        let source = String::from("del x");
+        let parse_ast = parse_program(&source, "<test>").unwrap();
+        insta::assert_debug_snapshot!(parse_ast);
+    }
+
+    #[test]
+    fn test_del_attribute() {
+        let source = String::from("del x.y");
+        let parse_ast = parse_program(&source, "<test>").unwrap();
+        insta::assert_debug_snapshot!(parse_ast);
+    }
+
+    #[test]
+    fn test_del_subscript() {
+        let source = String::from("del x[y]");
+        let parse_ast = parse_program(&source, "<test>").unwrap();
+        insta::assert_debug_snapshot!(parse_ast);
+    }
+}
diff --git a/parser/src/lib.rs b/parser/src/lib.rs
@@ -30,5 +30,6 @@ pub mod mode;
 pub mod parser;
 #[rustfmt::skip]
 mod python;
+mod context;
 mod string;
 pub mod token;
diff --git a/parser/src/snapshots/rustpython_parser__context__tests__ann_assign_name.snap b/parser/src/snapshots/rustpython_parser__context__tests__ann_assign_name.snap
@@ -4,40 +4,46 @@ expression: parse_ast
 ---
 [
     Located {
-        start: Location {
+        location: Location {
             row: 1,
             column: 1,
         },
-        end: Location {
-            row: 1,
-            column: 11,
-        },
+        end_location: Some(
+            Location {
+                row: 1,
+                column: 11,
+            },
+        ),
         custom: (),
         node: AnnAssign {
             target: Located {
-                start: Location {
+                location: Location {
                     row: 1,
                     column: 1,
                 },
-                end: Location {
-                    row: 1,
-                    column: 2,
-                },
+                end_location: Some(
+                    Location {
+                        row: 1,
+                        column: 2,
+                    },
+                ),
                 custom: (),
                 node: Name {
                     id: "x",
                     ctx: Store,
                 },
             },
             annotation: Located {
-                start: Location {
+                location: Location {
                     row: 1,
                     column: 4,
                 },
-                end: Location {
-                    row: 1,
-                    column: 7,
-                },
+                end_location: Some(
+                    Location {
+                        row: 1,
+                        column: 7,
+                    },
+                ),
                 custom: (),
                 node: Name {
                     id: "int",
@@ -46,14 +52,16 @@ expression: parse_ast
             },
             value: Some(
                 Located {
-                    start: Location {
+                    location: Location {
                         row: 1,
                         column: 10,
                     },
-                    end: Location {
-                        row: 1,
-                        column: 11,
-                    },
+                    end_location: Some(
+                        Location {
+                            row: 1,
+                            column: 11,
+                        },
+                    ),
                     custom: (),
                     node: Constant {
                         value: Int(

diff --git a/parser/src/snapshots/rustpython_parser__context__tests__assign.snap b/parser/src/snapshots/rustpython_parser__context__tests__assign.snap
@@ -0,0 +1,42 @@
+---
+source: compiler/parser/src/context.rs
+expression: parse_ast
+---
+[
+    Located {
+        location: Location {
+            row: 1,
+            column: 1,
+        },
+        custom: (),
+        node: Assign {
+            targets: [
+                Located {
+                    location: Location {
+                        row: 1,
+                        column: 1,
+                    },
+                    custom: (),
+                    node: Name {
+                        id: "x",
+                        ctx: Store,
+                    },
+                },
+            ],
+            value: Located {
+                location: Location {
+                    row: 1,
+                    column: 5,
+                },
+                custom: (),
+                node: Constant {
+                    value: Int(
+                        1,
+                    ),
+                    kind: None,
+                },
+            },
+            type_comment: None,
+        },
+    },
+]