Mercurial > python-compiler.rs

changeset 56:c3cc16b933d2

Find changesets by keywords (author, files, the commit message), revision number or hash, or revset expression.
Implement function arguments of all kinds.
author Emmanuel Gil Peyrot <linkmauve@linkmauve.fr>
date Sun, 12 Jun 2016 03:48:24 +0100
parents 32c78b275d5a
children e5a808ec31c0
files src/ast_convert.rs src/ast_dump.rs src/python_ast.rs tests/test_parse_files/test_functiondef.py
diffstat 4 files changed, 117 insertions(+), 19 deletions(-) [+]
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 {
--- a/src/ast_dump.rs
+++ b/src/ast_dump.rs
@@ -163,11 +163,46 @@ impl to_string_able for expr {
     }
 }
 
+impl to_string_able for arg {
+    fn to_string(&self) -> String {
+        match self.annotation {
+            None => self.arg.clone(),
+            Some(ref annotation) => format!("{}: {}", self.arg, annotation.to_string())
+        }
+    }
+}
+
 impl to_string_able for arguments {
     fn to_string(&self) -> String {
         let mut args = vec!();
-        for arg in self.args.clone() {
-            args.push(arg.arg);
+        let num_args = self.args.len();
+        let num_defaults = self.defaults.len();
+
+        let mut normal_args = self.args.clone();
+        let defaulted_args = normal_args.split_off(num_args - num_defaults);
+        for arg in normal_args {
+            args.push(arg.to_string());
+        }
+        for (arg, default) in defaulted_args.iter().zip(self.defaults.iter()) {
+            args.push(format!("{}={}", arg.to_string(), default.to_string()));
+        }
+
+        if !self.kwonlyargs.is_empty() {
+            match self.vararg {
+                None => args.push("*".to_string()),
+                Some(ref vararg) => args.push(format!("*{}", vararg.to_string())),
+            }
+        }
+        for (arg, default) in self.kwonlyargs.iter().zip(self.kw_defaults.iter()) {
+            let arg = arg.to_string();
+            match *default {
+                Some(ref default) => args.push(format!("{}={}", arg, default.to_string())),
+                None => args.push(arg)
+            }
+        }
+        match self.kwarg {
+            Some(ref kwarg) => args.push(format!("**{}", kwarg.to_string())),
+            None => ()
         }
         args.join(", ")
     }
--- a/src/python_ast.rs
+++ b/src/python_ast.rs
@@ -193,7 +193,7 @@ pub struct arguments {
     pub args: Vec<arg>,
     pub vararg: Option<arg>,
     pub kwonlyargs: Vec<arg>,
-    pub kw_defaults: Vec<expr>,
+    pub kw_defaults: Vec<Option<expr>>,
     pub kwarg: Option<arg>,
     pub defaults: Vec<expr>
 }
--- a/tests/test_parse_files/test_functiondef.py
+++ b/tests/test_parse_files/test_functiondef.py
@@ -1,2 +1,8 @@
-def func():
-    print(1)
+def func0():
+    pass
+def func1(a, b: str):
+    pass
+def func2(a, *, b: composite.type, c=2, d: int=3):
+    pass
+def func3(a, c=1, *args, b=2, **kwargs):
+    pass