changeset 16:b21a246349f3

Add ast.DictComp
author Bastien Orivel <eijebong@bananium.fr>
date Thu, 02 Jun 2016 22:20:48 +0200
parents a0fb169fe0f9
children f1c77634a2eb
files example/dictcomp.py src/ast_convert.rs src/ast_dump.rs src/python_ast.rs
diffstat 4 files changed, 22 insertions(+), 4 deletions(-) [+]
line wrap: on
line diff
new file mode 100644
--- /dev/null
+++ b/example/dictcomp.py
@@ -0,0 +1,1 @@
+d = {key:  4 for key in imaginary_list if key % 12}
--- a/src/ast_convert.rs
+++ b/src/ast_convert.rs
@@ -216,6 +216,7 @@ fn parse_expr(py: Python, ast: PyObject)
     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();
+    let dictcomp_type = ast_module.get(py, "DictComp").unwrap();
 
     assert!(is_instance(&ast, &ast_type));
 
@@ -335,8 +336,23 @@ fn parse_expr(py: Python, ast: PyObject)
             new_gens.push(gen);
         }
         expr::ListComp(Box::new(elt), new_gens)
-    }
-    else {
+    } else if is_instance(&ast, &dictcomp_type) {
+        let key = ast.getattr(py, "key").unwrap();
+        let value = ast.getattr(py, "value").unwrap();
+        let generators = ast.getattr(py, "generators").unwrap();
+        let generators = generators.iter(py).unwrap();
+
+        let key = parse_expr(py, key);
+        let value = parse_expr(py, value);
+
+        let mut new_gens = vec!();
+        for gen in generators {
+            let gen = gen.unwrap();
+            let gen = parse_comprehension(py, gen);
+            new_gens.push(gen);
+        }
+        expr::DictComp(Box::new(key), Box::new(value), new_gens)
+    } else {
         println!("expr {}", ast);
         unreachable!()
     }
--- a/src/ast_dump.rs
+++ b/src/ast_dump.rs
@@ -147,7 +147,8 @@ impl expr {
             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::ListComp(elt, generators) => format!("[{} {}]", elt.to_string(), generators_to_string(generators))
+            expr::ListComp(elt, generators) => format!("[{} {}]", elt.to_string(), generators_to_string(generators)),
+            expr::DictComp(key, value, generators) => format!("{{{}: {} {}}}", key.to_string(), value.to_string(), generators_to_string(generators))
         }
     }
 }
--- a/src/python_ast.rs
+++ b/src/python_ast.rs
@@ -86,7 +86,7 @@ pub enum expr {
     //Set(Vec<expr>)
     ListComp(Box<expr>, Vec<comprehension>),
     //SetComp(Box<expr>, Vec<comprehension>)
-    //DictComp(Box<expr>, Box<expr>, Vec<comprehension>)
+    DictComp(Box<expr>, Box<expr>, Vec<comprehension>),
     //GeneratorExp(Box<expr>, Vec<comprehension>)
 
     // the grammar constrains where yield expressions can occur