RustPython · cthulahoops · Mar 5, 2019 · Mar 4, 2019 · Mar 4, 2019
diff --git a/parser/src/ast.rs b/parser/src/ast.rs
@@ -122,6 +122,7 @@ pub enum Statement {
         // docstring: String,
         body: Vec<LocatedStatement>,
         decorator_list: Vec<Expression>,
+        returns: Option<Expression>,
     },
 }
 
@@ -269,14 +270,20 @@ impl Expression {
  */
 #[derive(Debug, PartialEq, Default)]
 pub struct Parameters {
-    pub args: Vec<String>,
-    pub kwonlyargs: Vec<String>,
-    pub vararg: Option<Option<String>>, // Optionally we handle optionally named '*args' or '*'
-    pub kwarg: Option<Option<String>>,
+    pub args: Vec<Parameter>,
+    pub kwonlyargs: Vec<Parameter>,
+    pub vararg: Option<Option<Parameter>>, // Optionally we handle optionally named '*args' or '*'
+    pub kwarg: Option<Option<Parameter>>,
     pub defaults: Vec<Expression>,
     pub kw_defaults: Vec<Option<Expression>>,
 }
 
+#[derive(Debug, PartialEq, Default)]
+pub struct Parameter {
+    pub arg: String,
+    pub annotation: Option<Box<Expression>>,
+}
+
 #[derive(Debug, PartialEq)]
 pub enum ComprehensionKind {
     GeneratorExpression { element: Expression },

diff --git a/parser/src/parser.rs b/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,
                             }
                         }
                     ],

diff --git a/parser/src/python.lalrpop b/parser/src/python.lalrpop
@@ -446,21 +446,22 @@ 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),
             }
         }
     },
 };
 
 Parameters: ast::Parameters = {
-    "(" <a: (TypedArgsList)?> ")" => {
+    "(" <a: (TypedArgsList<TypedParameter>)?> ")" => {
         match a {
             Some(a) => a,
             None => Default::default(),
@@ -470,8 +471,10 @@ Parameters: ast::Parameters = {
 
 // parameters are (String, None), kwargs are (String, Some(Test)) where Test is
 // the default
-TypedArgsList: ast::Parameters = {
-    <param1:TypedParameters> <args2:("," ParameterListStarArgs)?> => {
+// Note that this is a macro which is used once for function defs, and
+// once for lambda defs.
+TypedArgsList<ArgType>: ast::Parameters = {
+    <param1:TypedParameters<ArgType>> <args2:("," ParameterListStarArgs<ArgType>)?> => {
         let (names, default_elements) = param1;
 
         // Now gather rest of parameters:
@@ -489,7 +492,7 @@ TypedArgsList: ast::Parameters = {
             kw_defaults: kw_defaults,
         }
     },
-    <param1:TypedParameters> <kw:("," KwargParameter)> => {
+    <param1:TypedParameters<ArgType>> <kw:("," KwargParameter<ArgType>)> => {
         let (names, default_elements) = param1;
 
         // Now gather rest of parameters:
@@ -507,7 +510,7 @@ TypedArgsList: ast::Parameters = {
             kw_defaults: kw_defaults,
         }
     },
-    <params:ParameterListStarArgs> => {
+    <params:ParameterListStarArgs<ArgType>> => {
         let (vararg, kwonlyargs, kw_defaults, kwarg) = params;
         ast::Parameters {
             args: vec![],
@@ -518,7 +521,7 @@ TypedArgsList: ast::Parameters = {
             kw_defaults: kw_defaults,
         }
     },
-    <kw:KwargParameter> => {
+    <kw:KwargParameter<ArgType>> => {
         ast::Parameters {
             args: vec![],
             kwonlyargs: vec![],
@@ -532,8 +535,8 @@ TypedArgsList: ast::Parameters = {
 
 // Use inline here to make sure the "," is not creating an ambiguity.
 #[inline]
-TypedParameters: (Vec<String>, Vec<ast::Expression>) = {
-    <param1:TypedParameterDef> <param2:("," TypedParameterDef)*> => {
+TypedParameters<ArgType>: (Vec<ast::Parameter>, Vec<ast::Expression>) = {
+    <param1:TypedParameterDef<ArgType>> <param2:("," TypedParameterDef<ArgType>)*> => {
         // Combine first parameters:
         let mut args = vec![param1];
         args.extend(param2.into_iter().map(|x| x.1));
@@ -542,7 +545,6 @@ TypedParameters: (Vec<String>, Vec<ast::Expression>) = {
         let mut default_elements = vec![];
 
         for (name, default) in args.into_iter() {
-            names.push(name.clone());
             if let Some(default) = default {
                 default_elements.push(default);
             } else {
@@ -551,28 +553,35 @@ TypedParameters: (Vec<String>, Vec<ast::Expression>) = {
                     // have defaults
                     panic!(
                         "non-default argument follows default argument: {}",
-                        name
+                        &name.arg
                     );
                 }
             }
+            names.push(name);
         }
 
         (names, default_elements)
     }
 };
 
-TypedParameterDef: (String, Option<ast::Expression>) = {
-    <i:TypedParameter> => (i, None),
-    <i:TypedParameter> "=" <e:Test> => (i, Some(e)),
+TypedParameterDef<ArgType>: (ast::Parameter, Option<ast::Expression>) = {
+    <i:ArgType> => (i, None),
+    <i:ArgType> "=" <e:Test> => (i, Some(e)),
 };
 
-// TODO: add type annotations here:
-TypedParameter: String = {
-    Identifier,
+UntypedParameter: ast::Parameter = {
+    <i:Identifier> => ast::Parameter { arg: i, annotation: None },
 };
 
-ParameterListStarArgs: (Option<Option<String>>, Vec<String>, Vec<Option<ast::Expression>>, Option<Option<String>>) = {
-    "*" <va:Identifier?> <kw:("," TypedParameterDef)*> <kwarg:("," KwargParameter)?> => {
+TypedParameter: ast::Parameter = {
+    <arg:Identifier> <a:(":" Test)?>=> {
+        let annotation = a.map(|x| Box::new(x.1));
+        ast::Parameter { arg, annotation }
+    },
+};
+
+ParameterListStarArgs<ArgType>: (Option<Option<ast::Parameter>>, Vec<ast::Parameter>, Vec<Option<ast::Expression>>, Option<Option<ast::Parameter>>) = {
+    "*" <va:ArgType?> <kw:("," TypedParameterDef<ArgType>)*> <kwarg:("," KwargParameter<ArgType>)?> => {
         // Extract keyword arguments:
         let mut kwonlyargs = vec![];
         let mut kw_defaults = vec![];
@@ -590,8 +599,8 @@ ParameterListStarArgs: (Option<Option<String>>, Vec<String>, Vec<Option<ast::Exp
     }
 };
 
-KwargParameter: Option<String> = {
-    "**" <kwarg:Identifier?> => {
+KwargParameter<ArgType>: Option<ast::Parameter> = {
+    "**" <kwarg:ArgType?> => {
         kwarg
     }
 };
@@ -675,7 +684,7 @@ Test: ast::Expression = {
 };
 
 LambdaDef: ast::Expression = {
-    "lambda" <p:TypedArgsList?> ":" <b:Test> =>
+    "lambda" <p:TypedArgsList<UntypedParameter>?> ":" <b:Test> =>
         ast::Expression::Lambda {
             args: p.unwrap_or(Default::default()),
             body:Box::new(b)
@@ -1086,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,

diff --git a/tests/snippets/type_hints.py b/tests/snippets/type_hints.py
@@ -0,0 +1,7 @@
+
+# See also: https://github.com/RustPython/RustPython/issues/587
+
+def curry(foo: int) -> float:
+    return foo * 3.1415926 * 2
+
+assert curry(2) > 10
diff --git a/vm/src/compile.rs b/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,
@@ -442,10 +443,14 @@ impl Compiler {
 
         let line_number = self.get_source_line_number();
         self.code_object_stack.push(CodeObject::new(
-            args.args.clone(),
-            args.vararg.clone(),
-            args.kwonlyargs.clone(),
-            args.kwarg.clone(),
+            args.args.iter().map(|a| a.arg.clone()).collect(),
+            args.vararg
+                .as_ref()
+                .map(|x| x.as_ref().map(|a| a.arg.clone())),
+            args.kwonlyargs.iter().map(|a| a.arg.clone()).collect(),
+            args.kwarg
+                .as_ref()
+                .map(|x| x.as_ref().map(|a| a.arg.clone())),
             self.source_path.clone().unwrap(),
             line_number,
             name.to_string(),
@@ -584,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

diff --git a/vm/src/stdlib/ast.rs b/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.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");