Add return annotation and fix tests in parser.

Author: windelbouwman

parser/src/ast.rs

@@ -122,6 +122,7 @@ pub enum Statement {
         // docstring: String,
         body: Vec<LocatedStatement>,
         decorator_list: Vec<Expression>,
+        returns: Option<Expression>,
     },
 }
 

parser/src/parser.rs

@@ -244,7 +244,16 @@ mod tests {
                 node: ast::Statement::Expression {
                     expression: ast::Expression::Lambda {
                         args: ast::Parameters {
-                            args: vec![String::from("x"), String::from("y")],
+                            args: vec![
+                                ast::Parameter {
+                                    arg: String::from("x"),
+                                    annotation: None,
+                                },
+                                ast::Parameter {
+                                    arg: String::from("y"),
+                                    annotation: None,
+                                }
+                            ],
                             kwonlyargs: vec![],
                             vararg: None,
                             kwarg: None,
@@ -330,7 +339,10 @@ mod tests {
                             node: ast::Statement::FunctionDef {
                                 name: String::from("__init__"),
                                 args: ast::Parameters {
-                                    args: vec![String::from("self")],
+                                    args: vec![ast::Parameter {
+                                        arg: String::from("self"),
+                                        annotation: None,
+                                    }],
                                     kwonlyargs: vec![],
                                     vararg: None,
                                     kwarg: None,
@@ -342,14 +354,24 @@ mod tests {
                                     node: ast::Statement::Pass,
                                 }],
                                 decorator_list: vec![],
+                                returns: None,
                             }
                         },
                         ast::LocatedStatement {
                             location: ast::Location::new(4, 2),
                             node: ast::Statement::FunctionDef {
                                 name: String::from("method_with_default"),
                                 args: ast::Parameters {
-                                    args: vec![String::from("self"), String::from("arg"),],
+                                    args: vec![
+                                        ast::Parameter {
+                                            arg: String::from("self"),
+                                            annotation: None,
+                                        },
+                                        ast::Parameter {
+                                            arg: String::from("arg"),
+                                            annotation: None,
+                                        }
+                                    ],
                                     kwonlyargs: vec![],
                                     vararg: None,
                                     kwarg: None,
@@ -365,6 +387,7 @@ mod tests {
                                     node: ast::Statement::Pass,
                                 }],
                                 decorator_list: vec![],
+                                returns: None,
                             }
                         }
                     ],

parser/src/python.lalrpop

@@ -446,14 +446,15 @@ WithItem: ast::WithItem = {
 };
 
 FuncDef: ast::LocatedStatement = {
-    <d:Decorator*> <loc:@L> "def" <i:Identifier> <a:Parameters> ":" <s:Suite> => {
+    <d:Decorator*> <loc:@L> "def" <i:Identifier> <a:Parameters> <r:("->" Test)?> ":" <s:Suite> => {
         ast::LocatedStatement {
             location: loc,
             node: ast::Statement::FunctionDef {
                 name: i,
                 args: a,
                 body: s,
                 decorator_list: d,
+                returns: r.map(|x| x.1),
             }
         }
     },
@@ -1094,6 +1095,7 @@ extern {
         "<=" => lexer::Tok::LessEqual,
         ">" => lexer::Tok::Greater,
         ">=" => lexer::Tok::GreaterEqual,
+        "->" => lexer::Tok::Rarrow,
         "and" => lexer::Tok::And,
         "as" => lexer::Tok::As,
         "assert" => lexer::Tok::Assert,

tests/snippets/type_hints.py

@@ -1,7 +1,7 @@
 
 # See also: https://github.com/RustPython/RustPython/issues/587
 
-def curry(foo: int):  # TODO: -> float:
+def curry(foo: int) -> float:
     return foo * 3.1415926 * 2
 
 assert curry(2) > 10

vm/src/compile.rs

@@ -300,7 +300,8 @@ impl Compiler {
                 args,
                 body,
                 decorator_list,
-            } => self.compile_function_def(name, args, body, decorator_list)?,
+                returns,
+            } => self.compile_function_def(name, args, body, decorator_list, returns)?,
             ast::Statement::ClassDef {
                 name,
                 body,
@@ -588,6 +589,7 @@ impl Compiler {
         args: &ast::Parameters,
         body: &[ast::LocatedStatement],
         decorator_list: &[ast::Expression],
+        _returns: &Option<ast::Expression>, // TODO: use type hint somehow..
     ) -> Result<(), CompileError> {
         // Create bytecode for this function:
         // remember to restore self.in_loop to the original after the function is compiled

vm/src/stdlib/ast.rs

@@ -82,6 +82,7 @@ fn statement_to_ast(ctx: &PyContext, statement: &ast::LocatedStatement) -> PyObj
             args,
             body,
             decorator_list,
+            returns,
         } => {
             let node = create_node(ctx, "FunctionDef");
 
@@ -96,6 +97,13 @@ fn statement_to_ast(ctx: &PyContext, statement: &ast::LocatedStatement) -> PyObj
 
             let py_decorator_list = expressions_to_ast(ctx, decorator_list);
             ctx.set_attr(&node, "decorator_list", py_decorator_list);
+
+            let py_returns = if let Some(hint) = returns {
+                expression_to_ast(ctx, hint)
+            } else {
+                ctx.none()
+            };
+            ctx.set_attr(&node, "returns", py_returns);
             node
         }
         ast::Statement::Continue => create_node(ctx, "Continue"),
@@ -537,17 +545,28 @@ fn parameters_to_ast(ctx: &PyContext, args: &ast::Parameters) -> PyObjectRef {
     ctx.set_attr(
         &node,
         "args",
-        ctx.new_list(
-            args.args
-                .iter()
-                .map(|a| ctx.new_str(a.arg.to_string()))
-                .collect(),
-        ),
+        ctx.new_list(args.args.iter().map(|a| parameter_to_ast(ctx, a)).collect()),
     );
 
     node
 }
 
+fn parameter_to_ast(ctx: &PyContext, parameter: &ast::Parameter) -> PyObjectRef {
+    let node = create_node(ctx, "arg");
+
+    let py_arg = ctx.new_str(parameter.arg.to_string());
+    ctx.set_attr(&node, "arg", py_arg);
+
+    let py_annotation = if let Some(annotation) = &parameter.annotation {
+        expression_to_ast(ctx, annotation)
+    } else {
+        ctx.none()
+    };
+    ctx.set_attr(&node, "annotation", py_annotation);
+
+    node
+}
+
 fn comprehension_to_ast(ctx: &PyContext, comprehension: &ast::Comprehension) -> PyObjectRef {
     let node = create_node(ctx, "comprehension");