changeset 14:719a27f1c1c7

Add ast.ListComp
author Bastien Orivel <eijebong@bananium.fr>
date Thu, 02 Jun 2016 20:40:24 +0200
parents 38b0d63697b1
children a0fb169fe0f9
files example/listcomp.py src/ast_convert.rs src/ast_dump.rs src/python_ast.rs
diffstat 4 files changed, 51 insertions(+), 5 deletions(-) [+]
line wrap: on
line diff
new file mode 100644
--- /dev/null
+++ b/example/listcomp.py
@@ -0,0 +1,1 @@
+l = [v for v in range(5) if v % 2]
--- 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!()
     }
--- a/src/ast_dump.rs
+++ b/src/ast_dump.rs
@@ -1,4 +1,4 @@
-use python_ast::{Module, stmt, expr, boolop, operator, unaryop, cmpop, arguments, arg, keyword};
+use python_ast::{Module, stmt, expr, boolop, operator, unaryop, cmpop, arguments, arg, keyword, comprehension};
 
 use std::iter;
 
@@ -103,6 +103,16 @@ fn if_else_statements_to_string(indent: 
     }
 }
 
+fn generators_to_string(generators: Vec<comprehension>) -> String {
+    let mut result  = vec!();
+    for generator in generators {
+        result.push(format!("for {} in {}", generator.target.to_string(), generator.iter.to_string()));
+        for if_ in generator.ifs {
+            result.push(format!("if {}", if_.to_string()))
+        }
+    }
+    result.join(" ")
+}
 impl expr {
     fn to_string(&self) -> String {
         match self.clone() {
@@ -136,7 +146,8 @@ impl expr {
             expr::NameConstant(name) => format!("{}", name),
             expr::Attribute(value, attr, ctx) => format!("{}.{}", value.to_string(), attr),
             expr::Name(name, ctx) => format!("{}", name),
-            expr::List(elts, ctx) => format!("[{}]", args_to_string(elts))
+            expr::List(elts, ctx) => format!("[{}]", args_to_string(elts)),
+            expr::ListComp(elt, generators) => format!("[{} {}]", elt.to_string(), generators_to_string(generators))
         }
     }
 }
--- a/src/python_ast.rs
+++ b/src/python_ast.rs
@@ -84,7 +84,7 @@ pub enum expr {
     //IfExp(Box<expr>, Box<expr>, Box<expr>)
     //Dict(Vec<expr>, Vec<expr>)
     //Set(Vec<expr>)
-    //ListComp(Box<expr>, Vec<comprehension>),
+    ListComp(Box<expr>, Vec<comprehension>),
     //SetComp(Box<expr>, Vec<comprehension>)
     //DictComp(Box<expr>, Box<expr>, Vec<comprehension>)
     //GeneratorExp(Box<expr>, Vec<comprehension>)