Treewalk: Support multiple expressions for nonlocal & global, refacto…

…r CPythonModule to remain unimported until requested

src/parser/mod.rs

@@ -10,7 +10,7 @@ use crate::parser::types::{
     ExceptionLiteral, Expr, ExprFormat, FStringPart, ForClause, FormatOption, HandledException,
     ImportPath, ImportedItem, LogicalOp, LoopIndex, ParsedArgDefinition, ParsedArgDefinitions,
     ParsedArgument, ParsedArguments, ParsedSliceParams, RegularImport, Statement, TypeNode,
-    UnaryOp,
+    UnaryOp, Variable,
 };
 use crate::treewalk::State;
 use crate::types::errors::ParserError;
@@ -177,7 +177,9 @@ impl Parser {
 
                 // We need this for the case of a trailing comma, which is most often used for a
                 // tuple with a single element.
-                if self.end_of_statement() {
+                //
+                // The [`Token::Assign`] is when this happens on the LHS.
+                if self.end_of_statement() || self.current_token == Token::Assign {
                     break;
                 }
                 items.push(self.parse_simple_expr()?);
@@ -1081,7 +1083,7 @@ impl Parser {
         self.consume(Token::Colon)?;
         let if_part = ConditionalBlock {
             condition,
-            block: self.parse_indented_block()?,
+            block: self.parse_block()?,
         };
 
         let mut elif_parts: Vec<ConditionalBlock> = vec![];
@@ -1146,7 +1148,7 @@ impl Parser {
         })
     }
 
-    fn parse_body(&mut self) -> Result<Block, ParserError> {
+    fn parse_block(&mut self) -> Result<Block, ParserError> {
         if self.current_token == Token::Newline {
             self.parse_indented_block()
         } else {
@@ -1212,7 +1214,7 @@ impl Parser {
         }
 
         self.consume(Token::Colon)?;
-        let body = self.parse_body()?;
+        let body = self.parse_block()?;
 
         Ok(Statement::ClassDef {
             name,
@@ -1253,7 +1255,7 @@ impl Parser {
         }
 
         self.consume(Token::Colon)?;
-        let body = self.parse_body()?;
+        let body = self.parse_block()?;
 
         Ok(Statement::FunctionDef {
             name,
@@ -1335,11 +1337,13 @@ impl Parser {
             }
             Token::Nonlocal => {
                 self.consume(Token::Nonlocal)?;
-                Ok(Statement::Nonlocal(self.parse_identifier()?))
+                let identifiers = self.parse_identifiers()?;
+                Ok(Statement::Nonlocal(identifiers))
             }
             Token::Global => {
                 self.consume(Token::Global)?;
-                Ok(Statement::Global(self.parse_identifier()?))
+                let identifiers = self.parse_identifiers()?;
+                Ok(Statement::Global(identifiers))
             }
             Token::If => self.parse_if_else(),
             Token::While => {
@@ -1870,6 +1874,15 @@ impl Parser {
         }
     }
 
+    fn parse_identifiers(&mut self) -> Result<Vec<Variable>, ParserError> {
+        let mut items = vec![self.parse_identifier()?];
+        while self.current_token == Token::Comma {
+            self.consume(Token::Comma)?;
+            items.push(self.parse_identifier()?);
+        }
+        Ok(items)
+    }
+
     /// Parse a `Token::Identifier` without any semantic analysis.
     fn parse_identifier(&mut self) -> Result<String, ParserError> {
         match self.current_token.clone() {
@@ -2786,6 +2799,27 @@ if x > 0:
             Ok(ast) => assert_eq!(ast, expected_ast),
         }
 
+        let input = r#"
+if True: return False
+"#;
+        let mut parser = init(input);
+
+        let expected_ast = Statement::IfElse {
+            if_part: ConditionalBlock {
+                condition: Expr::Boolean(true),
+                block: Block {
+                    statements: vec![Statement::Return(vec![Expr::Boolean(false)])],
+                },
+            },
+            elif_parts: vec![],
+            else_part: None,
+        };
+
+        match parse(&mut parser) {
+            Err(e) => panic!("Parser error: {:?}", e),
+            Ok(ast) => assert_eq!(ast, expected_ast),
+        }
+
         let input = r#"
 if (a == 1
         and b
@@ -5662,7 +5696,19 @@ nonlocal var
 ";
         let mut parser = init(input);
 
-        let expected_ast = Statement::Nonlocal("var".into());
+        let expected_ast = Statement::Nonlocal(vec!["var".into()]);
+
+        match parse(&mut parser) {
+            Err(e) => panic!("Parser error: {:?}", e),
+            Ok(ast) => assert_eq!(ast, expected_ast),
+        }
+
+        let input = "
+nonlocal var, var2
+";
+        let mut parser = init(input);
+
+        let expected_ast = Statement::Nonlocal(vec!["var".into(), "var2".into()]);
 
         match parse(&mut parser) {
             Err(e) => panic!("Parser error: {:?}", e),
@@ -5674,7 +5720,19 @@ global var
 ";
         let mut parser = init(input);
 
-        let expected_ast = Statement::Global("var".into());
+        let expected_ast = Statement::Global(vec!["var".into()]);
+
+        match parse(&mut parser) {
+            Err(e) => panic!("Parser error: {:?}", e),
+            Ok(ast) => assert_eq!(ast, expected_ast),
+        }
+
+        let input = "
+global var, var2
+";
+        let mut parser = init(input);
+
+        let expected_ast = Statement::Global(vec!["var".into(), "var2".into()]);
 
         match parse(&mut parser) {
             Err(e) => panic!("Parser error: {:?}", e),
@@ -5910,6 +5968,21 @@ if True:
             Err(e) => panic!("Parser error: {:?}", e),
             Ok(ast) => assert_eq!(ast, expected_ast),
         }
+
+        let input = r#"
+a, = b,
+"#;
+        let mut parser = init(input);
+
+        let expected_ast = Statement::UnpackingAssignment {
+            left: vec![Expr::Variable("a".into())],
+            right: Expr::Tuple(vec![Expr::Variable("b".into())]),
+        };
+
+        match parse(&mut parser) {
+            Err(e) => panic!("Parser error: {:?}", e),
+            Ok(ast) => assert_eq!(ast, expected_ast),
+        }
     }
 
     #[test]

src/parser/static_analysis.rs

@@ -116,11 +116,10 @@ impl FunctionAnalysisVisitor {
                     }
                 }
             }
-            Statement::Nonlocal(var) => {
-                self.accessed_vars.push(var.into());
-            }
-            Statement::Global(var) => {
-                self.accessed_vars.push(var.into());
+            Statement::Global(names) | Statement::Nonlocal(names) => {
+                for name in names {
+                    self.accessed_vars.push(name.clone());
+                }
             }
             _ => {}
         }

src/parser/types.rs

@@ -1,4 +1,5 @@
 use std::collections::{HashMap, HashSet};
+use std::fmt::{Display, Error, Formatter};
 use std::hash::{Hash, Hasher};
 
 use crate::parser::static_analysis::{FunctionAnalysisVisitor, Visitor};
@@ -50,11 +51,11 @@ pub struct RegularImport {
     pub alias: Option<String>,
 }
 
-#[derive(Debug, PartialEq, Clone)]
+#[derive(Debug, PartialEq, Clone, Eq, Hash)]
 pub enum ImportPath {
     Absolute(Vec<String>),
     /// Any import path which _starts_ with a dot. The `usize` represents the number of levels up
-    /// to look for the path. This means that this exxample is in the same directory
+    /// to look for the path. This means that this example is in the same directory.
     ///
     /// ```python
     /// from .package.module import symbol
@@ -68,6 +69,12 @@ pub enum ImportPath {
     Relative(usize, Vec<String>),
 }
 
+impl Display for ImportPath {
+    fn fmt(&self, f: &mut Formatter<'_>) -> Result<(), Error> {
+        write!(f, "{}", self.as_str())
+    }
+}
+
 impl ImportPath {
     pub fn as_str(&self) -> String {
         match self {
@@ -585,8 +592,8 @@ pub enum Statement {
         body: Block,
     },
     Return(Vec<Expr>),
-    Nonlocal(Variable),
-    Global(Variable),
+    Nonlocal(Vec<Variable>),
+    Global(Vec<Variable>),
     IfElse {
         if_part: ConditionalBlock,
         elif_parts: Vec<ConditionalBlock>,

src/treewalk/call_stack.rs

@@ -37,6 +37,27 @@ impl StackFrame {
         }
     }
 
+    fn empty_path() -> String {
+        "<stdin>".into()
+    }
+
+    fn file_path_str(&self) -> String {
+        match &self.file_path {
+            Some(path) => path.to_str().unwrap_or(&Self::empty_path()).to_string(),
+            None => Self::empty_path(),
+        }
+    }
+
+    fn empty_function_name() -> String {
+        "<module>".into()
+    }
+
+    fn function_name(&self) -> String {
+        self.function_name
+            .clone()
+            .unwrap_or(Self::empty_function_name())
+    }
+
     fn set_line(&mut self, line: usize) {
         self.line_number = line;
     }
@@ -102,19 +123,12 @@ impl Display for CallStack {
     fn fmt(&self, f: &mut Formatter) -> Result<(), Error> {
         writeln!(f, "Traceback (most recent call last):")?;
         for frame in &self.frames {
-            let file_path = frame
-                .file_path
-                .as_ref()
-                .and_then(|f| f.to_str())
-                .unwrap_or("<stdin>");
-            let context = frame
-                .function_name
-                .clone()
-                .unwrap_or("<module>".to_string());
             writeln!(
                 f,
                 "  File \"{}\", line {}, in {}",
-                file_path, frame.line_number, context
+                frame.file_path_str(),
+                frame.line_number,
+                frame.function_name()
             )?;
         }
         Ok(())