test a simple if statement

Author: jvllmr

Cargo.toml

@@ -8,3 +8,4 @@ rustpython-ast = { version = "0.4.0" }
 
 [dev-dependencies]
 rustpython-parser = "0.4.0"
+rustpython-ast = { version = "0.4.0", features = ["fold"] }

src/lib.rs

@@ -5,12 +5,39 @@ pub use crate::unparser::Unparser;
 #[cfg(test)]
 mod tests {
     use super::*;
+    use rustpython_ast::text_size::TextRange;
+    use rustpython_ast::Fold;
+
+    use rustpython_ast::TextSize;
     use rustpython_parser::ast::Suite;
     use rustpython_parser::Parse;
 
     use std::fs;
     use std::io;
     use std::path::Path;
+
+    struct RangesEraser {}
+
+    impl Fold<TextRange> for RangesEraser {
+        type TargetU = TextRange;
+
+        type Error = std::convert::Infallible;
+
+        type UserContext = TextRange;
+
+        fn will_map_user(&mut self, _user: &TextRange) -> Self::UserContext {
+            TextRange::new(TextSize::new(0), TextSize::new(0))
+        }
+
+        fn map_user(
+            &mut self,
+            _user: TextRange,
+            start: Self::UserContext,
+        ) -> Result<Self::TargetU, Self::Error> {
+            Ok(start)
+        }
+    }
+
     #[test]
     fn test_predefined_files() -> io::Result<()> {
         for entry in fs::read_dir("./test_files")? {
@@ -34,7 +61,19 @@ mod tests {
                 fs::write(&new_entry_path, &new_source)?;
                 let new_stmts =
                     Suite::parse(&new_source, new_entry_path.to_str().unwrap()).unwrap();
-                for (stmt, new_stmt) in stmts.iter().zip(new_stmts.iter()) {
+                // erase range information
+                let mut eraser = RangesEraser {};
+                let mut erased_new_stmts = Vec::new();
+                for stmt in &new_stmts {
+                    erased_new_stmts.push(eraser.fold_stmt(stmt.to_owned()).unwrap());
+                }
+
+                let mut erased_stmts = Vec::new();
+                for stmt in &stmts {
+                    erased_stmts.push(eraser.fold_stmt(stmt.to_owned()).unwrap());
+                }
+
+                for (stmt, new_stmt) in erased_stmts.iter().zip(erased_new_stmts.iter()) {
                     assert_eq!(stmt, new_stmt)
                 }
             }

src/unparser.rs

@@ -858,6 +858,7 @@ impl Unparser {
         self.unparse_expr(&node.func);
         let mut args_iter = node.args.iter().peekable();
         let mut keywords_iter = node.keywords.iter().peekable();
+        self.write_str("(");
         while let Some(arg) = args_iter.next() {
             self.unparse_expr(arg);
             if args_iter.peek().is_some() || keywords_iter.peek().is_some() {
@@ -870,6 +871,7 @@ impl Unparser {
                 self.write_str(", ");
             }
         }
+        self.write_str(")");
     }
 
     fn unparse_expr_formatted_value(&mut self, node: &ExprFormattedValue<TextRange>) {
@@ -915,7 +917,7 @@ impl Unparser {
                 self.write_str(&utf8.unwrap());
             }
             Constant::Int(value) => self.write_str(&value.to_string()),
-            Constant::Str(value) => self.write_str(value),
+            Constant::Str(value) => self.write_str(&format!("\"{}\"", value)),
             Constant::None => self.write_str("None"),
             Constant::Complex { real, imag: _ } => self.write_str(&real.to_string()),
             Constant::Float(value) => self.write_str(&value.to_string()),

test_files/simple_if_statement.py

@@ -0,0 +1,5 @@
+if __name__ == "__main__":
+    pass
+
+if True:
+    print("Hello!")