changeset 5:ddf372373a77

Add ast.For, ast.UnaryOp, and Sub and Div to ast.BinOp.
author Emmanuel Gil Peyrot <linkmauve@linkmauve.fr>
date Tue, 31 May 2016 04:15:00 +0100
parents f27a4aee9dfa
children 6f2bf13f4cb5
files src/ast_convert.rs src/ast_dump.rs src/ast_rewrite.rs src/ast_type.rs src/python_ast.rs src/python_dump.rs
diffstat 6 files changed, 156 insertions(+), 4 deletions(-) [+]
line wrap: on
line diff
--- a/src/ast_convert.rs
+++ b/src/ast_convert.rs
@@ -1,4 +1,4 @@
-use python_ast::{Module, Statement, Expr, BinOp};
+use python_ast::{Module, Statement, Expr, BinOp, UnaryOp};
 
 use cpython::{Python, PyObject};
 use cpython::ObjectProtocol; //for call method
@@ -36,6 +36,31 @@ fn parse_expr_vec(py: Python, ast: PyObj
     }
 }
 
+fn parse_unaryop(py: Python, ast: PyObject) -> UnaryOp {
+    let builtins_module = py.import("builtins").unwrap();
+    let isinstance = builtins_module.get(py, "isinstance").unwrap();
+
+    let is_instance = |object: &PyObject, type_: &PyObject| {
+        return isinstance.call(py, (object, type_), None).unwrap().is_true(py).unwrap();
+    };
+
+    let ast_module = py.import("ast").unwrap();
+    let ast_type = ast_module.get(py, "AST").unwrap();
+    let uadd_type = ast_module.get(py, "UAdd").unwrap();
+    let usub_type = ast_module.get(py, "USub").unwrap();
+
+    assert!(is_instance(&ast, &ast_type));
+
+    if is_instance(&ast, &uadd_type) {
+        UnaryOp::UAdd
+    } else if is_instance(&ast, &usub_type) {
+        UnaryOp::USub
+    } else {
+        println!("UnaryOp {}", ast);
+        UnaryOp::Error
+    }
+}
+
 fn parse_binop(py: Python, ast: PyObject) -> BinOp {
     let builtins_module = py.import("builtins").unwrap();
     let isinstance = builtins_module.get(py, "isinstance").unwrap();
@@ -50,6 +75,8 @@ fn parse_binop(py: Python, ast: PyObject
     let mult_type = ast_module.get(py, "Mult").unwrap();
     let eq_type = ast_module.get(py, "Eq").unwrap();
     let lt_type = ast_module.get(py, "Lt").unwrap();
+    let sub_type = ast_module.get(py, "Sub").unwrap();
+    let div_type = ast_module.get(py, "Div").unwrap();
 
     assert!(is_instance(&ast, &ast_type));
 
@@ -61,6 +88,10 @@ fn parse_binop(py: Python, ast: PyObject
         BinOp::BinEq
     } else if is_instance(&ast, &lt_type) {
         BinOp::BinLt
+    } else if is_instance(&ast, &sub_type) {
+        BinOp::Sub
+    } else if is_instance(&ast, &div_type) {
+        BinOp::Div
     } else {
         println!("BinOp {}", ast);
         BinOp::Error
@@ -78,6 +109,7 @@ fn parse_expr(py: Python, ast: PyObject)
     let ast_module = py.import("ast").unwrap();
     let ast_type = ast_module.get(py, "AST").unwrap();
     let arg_type = ast_module.get(py, "arg").unwrap();
+    let unary_op_type = ast_module.get(py, "UnaryOp").unwrap();
     let bin_op_type = ast_module.get(py, "BinOp").unwrap();
     let name_constant_type = ast_module.get(py, "NameConstant").unwrap();
     let attribute_type = ast_module.get(py, "Attribute").unwrap();
@@ -118,6 +150,14 @@ fn parse_expr(py: Python, ast: PyObject)
         let s = ast.getattr(py, "s").unwrap();
         let s = get_str(py, s);
         Expr::Str(s)
+    } else if is_instance(&ast, &unary_op_type) {
+        let op = ast.getattr(py, "op").unwrap();
+        let operand = ast.getattr(py, "operand").unwrap();
+
+        let op = parse_unaryop(py, op);
+        let operand = parse_expr(py, operand);
+
+        Expr::UnaryOp(op, Box::new(operand))
     } else if is_instance(&ast, &bin_op_type) {
         let left = ast.getattr(py, "left").unwrap();
         let op = ast.getattr(py, "op").unwrap();
@@ -204,6 +244,7 @@ fn parse_statement(py: Python, ast: PyOb
     let return_type = ast_module.get(py, "Return").unwrap();
     let import_from_type = ast_module.get(py, "ImportFrom").unwrap();
     let if_type = ast_module.get(py, "If").unwrap();
+    let for_type = ast_module.get(py, "For").unwrap();
     let expr_type = ast_module.get(py, "Expr").unwrap();
 
     assert!(is_instance(&ast, &ast_type));
@@ -301,6 +342,30 @@ fn parse_statement(py: Python, ast: PyOb
         }
 
         Statement::If(test, statements, orelse_)
+    } else if is_instance(&ast, &for_type) {
+        let target = ast.getattr(py, "target").unwrap();
+        let iter = ast.getattr(py, "iter").unwrap();
+        let body = ast.getattr(py, "body").unwrap();
+        let orelse = ast.getattr(py, "orelse").unwrap();
+
+        let target = parse_expr(py, target);
+        let iter = parse_expr(py, iter);
+
+        let mut statements = vec!();
+        for statement in body.iter(py).unwrap() {
+            let statement = statement.unwrap();
+            let statement = parse_statement(py, statement);
+            statements.push(statement);
+        }
+
+        let mut orelse_ = vec!();
+        for statement in orelse.iter(py).unwrap() {
+            let statement = statement.unwrap();
+            let statement = parse_statement(py, statement);
+            orelse_.push(statement);
+        }
+
+        Statement::For(target, iter, statements, orelse_)
     } else if is_instance(&ast, &assign_type) {
         let targets = ast.getattr(py, "targets").unwrap();
         let value = ast.getattr(py, "value").unwrap();
--- a/src/ast_dump.rs
+++ b/src/ast_dump.rs
@@ -1,7 +1,17 @@
-use python_ast::{Module, Statement, Expr, BinOp};
+use python_ast::{Module, Statement, Expr, BinOp, UnaryOp};
 
 use std::iter;
 
+impl UnaryOp {
+    fn to_string(&self) -> &'static str {
+        match *self {
+            UnaryOp::UAdd => "+",
+            UnaryOp::USub => "-",
+            UnaryOp::Error => "BinOp::Error"
+        }
+    }
+}
+
 impl BinOp {
     fn to_string(&self) -> &'static str {
         match *self {
@@ -9,6 +19,8 @@ impl BinOp {
             BinOp::BinMult => "*",
             BinOp::BinEq => "==",
             BinOp::BinLt => "<",
+            BinOp::Sub => "-",
+            BinOp::Div => "/",
             BinOp::Error => "BinOp::Error"
         }
     }
@@ -17,6 +29,7 @@ impl BinOp {
 impl Expr {
     fn to_string(&self) -> String {
         match self.clone() {
+            Expr::UnaryOp(op, operand) => format!("{}{}", op.to_string(), operand.to_string()),
             Expr::BinOp(a, op, b) => format!("{} {} {}", a.to_string(), op.to_string(), b.to_string()),
             Expr::Compare(left, ops, comparators) => format!("{} {}", left.to_string(), {
                 let mut arguments = vec!();
@@ -113,6 +126,23 @@ impl Statement {
                     format!("{}\n{}else:\n{}", statements.join("\n"), make_indent(indent), orelse_.join("\n"))
                 }
             }),
+            Statement::For(target, iter, body, orelse) => format!("{}for {} in {}:\n{}", make_indent(indent), target.to_string(), iter.to_string(), {
+                let mut statements = vec!();
+                for arg in body {
+                    statements.push(arg.to_string(indent + 1));
+                }
+
+                let mut orelse_ = vec!();
+                for arg in orelse {
+                    orelse_.push(arg.to_string(indent + 1));
+                }
+
+                if orelse_.is_empty() {
+                    statements.join("\n")
+                } else {
+                    format!("{}\n{}else:\n{}", statements.join("\n"), make_indent(indent), orelse_.join("\n"))
+                }
+            }),
             Statement::Assign(targets, value) => format!("{}{} = {}", make_indent(indent), {
                 let mut exprs = vec!();
                 for target in targets {
--- a/src/ast_rewrite.rs
+++ b/src/ast_rewrite.rs
@@ -40,6 +40,16 @@ impl Visitor<()> for Rewrite {
                     self.visit_statement(statement);
                 }
             },
+            Statement::For(target, iter, body, orelse) => {
+                self.visit_expr(target);
+                self.visit_expr(iter);
+                for statement in body {
+                    self.visit_statement(statement);
+                }
+                for statement in orelse {
+                    self.visit_statement(statement);
+                }
+            },
             Statement::Assign(targets, value) => {
                 self.visit_expr(value);
                 for target in targets {
--- a/src/ast_type.rs
+++ b/src/ast_type.rs
@@ -84,6 +84,17 @@ impl Visitor<Type> for Typing {
                 }
                 Type::Bottom
             },
+            Statement::For(target, iter, body, orelse) => {
+                self.visit_expr(target);
+                self.visit_expr(iter);
+                for statement in body {
+                    self.visit_statement(statement);
+                }
+                for statement in orelse {
+                    self.visit_statement(statement);
+                }
+                Type::Bottom
+            },
             Statement::Assign(targets, value) => {
                 let type_ = self.visit_expr(value);
                 if targets.len() != 1 {
@@ -125,6 +136,10 @@ impl Visitor<Type> for Typing {
     fn visit_expr(&mut self, expr: Expr) -> Type {
         let expr_str = format!("{:?}", expr);
         let type_ = match expr {
+            Expr::UnaryOp(op, operand) => {
+                let type_operand = self.visit_expr(*operand);
+                type_operand
+            },
             Expr::BinOp(left, op, right) => {
                 let type_left = self.visit_expr(*left);
                 let type_right = self.visit_expr(*right);
--- a/src/python_ast.rs
+++ b/src/python_ast.rs
@@ -10,15 +10,17 @@ pub enum Statement {
     FunctionDef(String, Vec<Expr>, Vec<Statement>),
     Global(Vec<String>),
     If(Expr, Vec<Statement>, Vec<Statement>),
+    For(Expr, Expr, Vec<Statement>, Vec<Statement>),
     Assign(Vec<Expr>, Expr),
     Return(Expr),
     ImportFrom(String, Vec<Expr>),
     Expr(Expr),
-    Error,
+    Error
 }
 
 #[derive(Clone, Debug, PartialEq, Eq, Hash)]
 pub enum Expr {
+    UnaryOp(UnaryOp, Box<Expr>),
     BinOp(Box<Expr>, BinOp, Box<Expr>),
     Compare(Box<Expr>, Vec<BinOp>, Vec<Expr>),
     Call(Box<Expr>, Vec<Expr>),
@@ -32,10 +34,19 @@ pub enum Expr {
 }
 
 #[derive(Clone, Debug, PartialEq, Eq, Hash)]
+pub enum UnaryOp {
+    UAdd,
+    USub,
+    Error
+}
+
+#[derive(Clone, Debug, PartialEq, Eq, Hash)]
 pub enum BinOp {
     BinAdd,
     BinMult,
     BinEq,
     BinLt,
-    Error,
+    Sub,
+    Div,
+    Error
 }
--- a/src/python_dump.rs
+++ b/src/python_dump.rs
@@ -30,6 +30,7 @@ fn dump(py: Python, indent: usize, ast: 
     let eq_type = ast_module.get(py, "Eq").unwrap();
     let lt_type = ast_module.get(py, "Lt").unwrap();
     let if_type = ast_module.get(py, "If").unwrap();
+    let for_type = ast_module.get(py, "For").unwrap();
     let compare_type = ast_module.get(py, "Compare").unwrap();
     let expr_type = ast_module.get(py, "Expr").unwrap();
     let call_type = ast_module.get(py, "Call").unwrap();
@@ -91,6 +92,26 @@ fn dump(py: Python, indent: usize, ast: 
         let indent = indent.as_str();
         let body = statements.join(indent);
         format!("{}{}", declaration, body)
+    } else if is_instance(&ast, &for_type) {
+        let target = ast.getattr(py, "target").unwrap();
+        let iter = ast.getattr(py, "iter").unwrap();
+        let body = ast.getattr(py, "body").unwrap();
+        //let orelse = ast.getattr(py, "orelse").unwrap();
+
+        let target = dump(py, indent, target);
+        let iter = dump(py, indent, iter);
+
+        let declaration = format!("for {} in {}:", target, iter);
+        let mut statements = vec!("".to_string());
+        for statement in body.iter(py).unwrap() {
+            let statement = dump(py, indent + 1, statement.unwrap());
+            statements.push(statement);
+        }
+        let indent: String = iter::repeat("    ").take(indent + 1).collect();
+        let indent = format!("\n{}", indent);
+        let indent = indent.as_str();
+        let body = statements.join(indent);
+        format!("{}{}", declaration, body)
     } else if is_instance(&ast, &arguments_type) {
         let args_list = ast.getattr(py, "args").unwrap();
         let mut arguments = vec!();