diff src/ast_convert.rs @ 56:c3cc16b933d2

Implement function arguments of all kinds.
author Emmanuel Gil Peyrot <linkmauve@linkmauve.fr>
date Sun, 12 Jun 2016 03:48:24 +0100
parents 1a815946c2e5
children e5a808ec31c0
line wrap: on
line diff
--- a/src/ast_convert.rs
+++ b/src/ast_convert.rs
@@ -251,7 +251,6 @@ 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 bool_op_type = ast_module.get(py, "BoolOp").unwrap();
     let bin_op_type = ast_module.get(py, "BinOp").unwrap();
@@ -270,11 +269,7 @@ fn parse_expr(py: Python, ast: PyObject)
 
     assert!(is_instance(&ast, &ast_type));
 
-    if is_instance(&ast, &arg_type) {
-        let arg = ast.getattr(py, "arg").unwrap();
-        let arg = get_str(py, arg);
-        expr::Name(arg, get_ctx(py, ast))
-    } else if is_instance(&ast, &attribute_type) {
+    if is_instance(&ast, &attribute_type) {
         let value = ast.getattr(py, "value").unwrap();
         let attr = ast.getattr(py, "attr").unwrap();
 
@@ -378,6 +373,20 @@ fn parse_expr(py: Python, ast: PyObject)
     }
 }
 
+fn parse_arg(py: Python, ast: PyObject) -> arg {
+    let arg = ast.getattr(py, "arg").unwrap();
+    let annotation = ast.getattr(py, "annotation").unwrap();
+
+    let arg = get_str(py, arg);
+    let annotation = if annotation == py.None() {
+        None
+    } else {
+        Some(parse_expr(py, annotation))
+    };
+
+    arg{arg: arg, annotation: annotation}
+}
+
 fn parse_arguments(py: Python, ast: PyObject) -> arguments {
     let builtins_module = py.import("builtins").unwrap();
     let isinstance = builtins_module.get(py, "isinstance").unwrap();
@@ -402,21 +411,69 @@ fn parse_arguments(py: Python, ast: PyOb
                 for arg in args.iter(py).unwrap() {
                     let arg = arg.unwrap();
                     assert!(is_instance(&arg, &arg_type));
-                    let arg = get_str(py, arg);
-                    arguments.push(arg{arg: arg, annotation: None});
+                    let arg = parse_arg(py, arg);
+                    arguments.push(arg);
                 }
                 arguments
             },
             //vararg: Option<arg>,
-            vararg: None,
+            vararg: {
+                let vararg = ast.getattr(py, "vararg").unwrap();
+                if vararg == py.None() {
+                    None
+                } else {
+                    let arg = parse_arg(py, vararg);
+                    Some(arg)
+                }
+            },
             //kwonlyargs: Vec<arg>,
-            kwonlyargs: vec!(),
-            //kw_defaults: Vec<expr>,
-            kw_defaults: vec!(),
+            kwonlyargs: {
+                let kwonlyargs = ast.getattr(py, "kwonlyargs").unwrap();
+                let mut arguments = vec!();
+                for arg in kwonlyargs.iter(py).unwrap() {
+                    let arg = arg.unwrap();
+                    assert!(is_instance(&arg, &arg_type));
+                    let arg = parse_arg(py, arg);
+                    arguments.push(arg);
+                }
+                arguments
+            },
+            //kw_defaults: Vec<Option<expr>>,
+            kw_defaults: {
+                let kw_defaults = ast.getattr(py, "kw_defaults").unwrap();
+                let mut arguments = vec!();
+                for arg in kw_defaults.iter(py).unwrap() {
+                    let arg = arg.unwrap();
+                    let arg = if arg == py.None() {
+                        None
+                    } else {
+                        Some(parse_expr(py, arg))
+                    };
+                    arguments.push(arg);
+                }
+                arguments
+            },
             //kwarg: Option<arg>,
-            kwarg: None,
+            kwarg: {
+                let kwarg = ast.getattr(py, "kwarg").unwrap();
+                if kwarg == py.None() {
+                    None
+                } else {
+                    let arg = parse_arg(py, kwarg);
+                    Some(arg)
+                }
+            },
             //defaults: Vec<expr>
-            defaults: vec!()
+            defaults: {
+                let defaults = ast.getattr(py, "defaults").unwrap();
+                let mut arguments = vec!();
+                for arg in defaults.iter(py).unwrap() {
+                    let arg = arg.unwrap();
+                    let arg = parse_expr(py, arg);
+                    arguments.push(arg);
+                }
+                arguments
+            }
         };
         args
     } else {