# HG changeset patch # User Bastien Orivel # Date 1464892824 -7200 # Node ID 719a27f1c1c7a1344fd5630cab06ea1433f6f6ee # Parent 38b0d63697b1e6b8b367814baa8bf88aed307dcd Add ast.ListComp diff --git a/example/listcomp.py b/example/listcomp.py new file mode 100644 --- /dev/null +++ b/example/listcomp.py @@ -0,0 +1,1 @@ +l = [v for v in range(5) if v % 2] diff --git a/src/ast_convert.rs b/src/ast_convert.rs --- 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!() } diff --git a/src/ast_dump.rs b/src/ast_dump.rs --- 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) -> 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)) } } } diff --git a/src/python_ast.rs b/src/python_ast.rs --- a/src/python_ast.rs +++ b/src/python_ast.rs @@ -84,7 +84,7 @@ pub enum expr { //IfExp(Box, Box, Box) //Dict(Vec, Vec) //Set(Vec) - //ListComp(Box, Vec), + ListComp(Box, Vec), //SetComp(Box, Vec) //DictComp(Box, Box, Vec) //GeneratorExp(Box, Vec)