diff src/ast_convert.rs @ 14:719a27f1c1c7

Add ast.ListComp
author Bastien Orivel <eijebong@bananium.fr>
date Thu, 02 Jun 2016 20:40:24 +0200
parents 38b0d63697b1
children a0fb169fe0f9
line wrap: on
line diff
--- a/src/ast_convert.rs
+++ b/src/ast_convert.rs
@@ -1,4 +1,4 @@
-use python_ast::{Module, stmt, expr, expr_context, cmpop, operator, unaryop, arguments, arg, alias};
+use python_ast::{Module, stmt, expr, expr_context, cmpop, operator, unaryop, arguments, arg, alias, comprehension};
 
 use cpython::{Python, PyObject};
 use cpython::ObjectProtocol; //for call method
@@ -100,6 +100,24 @@ fn parse_cmpop(py: Python, ast: PyObject
     }
 }
 
+fn parse_comprehension(py: Python, ast: PyObject) -> comprehension {
+        let target = ast.getattr(py, "target").unwrap();
+        let iter = ast.getattr(py, "iter").unwrap();
+        let ifs = ast.getattr(py, "ifs").unwrap();
+        let ifs = ifs.iter(py).unwrap();
+
+        let target = parse_expr(py, target);
+        let iter = parse_expr(py, iter);
+
+        let mut new_ifs = vec!();
+        for if_ in ifs {
+            let if_ = parse_expr(py, if_.unwrap());
+            new_ifs.push(if_);
+        }
+        comprehension {target: target, iter: iter, ifs: new_ifs}
+
+}
+
 fn parse_operator(py: Python, ast: PyObject) -> operator {
     let builtins_module = py.import("builtins").unwrap();
     let isinstance = builtins_module.get(py, "isinstance").unwrap();
@@ -179,6 +197,7 @@ fn parse_expr(py: Python, ast: PyObject)
     let list_type = ast_module.get(py, "List").unwrap();
     let compare_type = ast_module.get(py, "Compare").unwrap();
     let call_type = ast_module.get(py, "Call").unwrap();
+    let listcomp_type = ast_module.get(py, "ListComp").unwrap();
 
     assert!(is_instance(&ast, &ast_type));
 
@@ -284,7 +303,22 @@ fn parse_expr(py: Python, ast: PyObject)
         }
 
         expr::Compare(Box::new(left), new_ops, new_comparators)
-    } else {
+    } else if is_instance(&ast, &listcomp_type) {
+        let elt = ast.getattr(py, "elt").unwrap();
+        let generators = ast.getattr(py, "generators").unwrap();
+        let generators = generators.iter(py).unwrap();
+
+        let elt = parse_expr(py, elt);
+
+        let mut new_gens = vec!();
+        for gen in generators {
+            let gen = gen.unwrap();
+            let gen = parse_comprehension(py, gen);
+            new_gens.push(gen);
+        }
+        expr::ListComp(Box::new(elt), new_gens)
+    }
+    else {
         println!("expr {}", ast);
         unreachable!()
     }