Skip to content

Commit f7a2225

Browse files
windelbouwmanwindelbouwman
committed
Add return annotation and fix tests in parser.
1 parent 6544f60 commit f7a2225

6 files changed

Lines changed: 59 additions & 12 deletions

File tree

‎parser/src/ast.rs‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -122,6 +122,7 @@ pub enum Statement {
122122
// docstring: String,
123123
body: Vec<LocatedStatement>,
124124
decorator_list: Vec<Expression>,
125+
returns: Option<Expression>,
125126
},
126127
}
127128

‎parser/src/parser.rs‎

Lines changed: 26 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -244,7 +244,16 @@ mod tests {
244244
node: ast::Statement::Expression {
245245
expression: ast::Expression::Lambda {
246246
args: ast::Parameters {
247-
args: vec![String::from("x"), String::from("y")],
247+
args: vec![
248+
ast::Parameter {
249+
arg: String::from("x"),
250+
annotation: None,
251+
},
252+
ast::Parameter {
253+
arg: String::from("y"),
254+
annotation: None,
255+
}
256+
],
248257
kwonlyargs: vec![],
249258
vararg: None,
250259
kwarg: None,
@@ -330,7 +339,10 @@ mod tests {
330339
node: ast::Statement::FunctionDef {
331340
name: String::from("__init__"),
332341
args: ast::Parameters {
333-
args: vec![String::from("self")],
342+
args: vec![ast::Parameter {
343+
arg: String::from("self"),
344+
annotation: None,
345+
}],
334346
kwonlyargs: vec![],
335347
vararg: None,
336348
kwarg: None,
@@ -342,14 +354,24 @@ mod tests {
342354
node: ast::Statement::Pass,
343355
}],
344356
decorator_list: vec![],
357+
returns: None,
345358
}
346359
},
347360
ast::LocatedStatement {
348361
location: ast::Location::new(4, 2),
349362
node: ast::Statement::FunctionDef {
350363
name: String::from("method_with_default"),
351364
args: ast::Parameters {
352-
args: vec![String::from("self"), String::from("arg"),],
365+
args: vec![
366+
ast::Parameter {
367+
arg: String::from("self"),
368+
annotation: None,
369+
},
370+
ast::Parameter {
371+
arg: String::from("arg"),
372+
annotation: None,
373+
}
374+
],
353375
kwonlyargs: vec![],
354376
vararg: None,
355377
kwarg: None,
@@ -365,6 +387,7 @@ mod tests {
365387
node: ast::Statement::Pass,
366388
}],
367389
decorator_list: vec![],
390+
returns: None,
368391
}
369392
}
370393
],

‎parser/src/python.lalrpop‎

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -446,14 +446,15 @@ WithItem: ast::WithItem = {
446446
};
447447

448448
FuncDef: ast::LocatedStatement = {
449-
<d:Decorator*> <loc:@L> "def" <i:Identifier> <a:Parameters> ":" <s:Suite> => {
449+
<d:Decorator*> <loc:@L> "def" <i:Identifier> <a:Parameters> <r:("->" Test)?> ":" <s:Suite> => {
450450
ast::LocatedStatement {
451451
location: loc,
452452
node: ast::Statement::FunctionDef {
453453
name: i,
454454
args: a,
455455
body: s,
456456
decorator_list: d,
457+
returns: r.map(|x| x.1),
457458
}
458459
}
459460
},
@@ -1094,6 +1095,7 @@ extern {
10941095
"<=" => lexer::Tok::LessEqual,
10951096
">" => lexer::Tok::Greater,
10961097
">=" => lexer::Tok::GreaterEqual,
1098+
"->" => lexer::Tok::Rarrow,
10971099
"and" => lexer::Tok::And,
10981100
"as" => lexer::Tok::As,
10991101
"assert" => lexer::Tok::Assert,

‎tests/snippets/type_hints.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,7 @@
11

22
# See also: https://github.com/RustPython/RustPython/issues/587
33

4-
def curry(foo: int): # TODO: -> float:
4+
def curry(foo: int) -> float:
55
return foo * 3.1415926 * 2
66

77
assert curry(2) > 10

‎vm/src/compile.rs‎

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -300,7 +300,8 @@ impl Compiler {
300300
args,
301301
body,
302302
decorator_list,
303-
} => self.compile_function_def(name, args, body, decorator_list)?,
303+
returns,
304+
} => self.compile_function_def(name, args, body, decorator_list, returns)?,
304305
ast::Statement::ClassDef {
305306
name,
306307
body,
@@ -588,6 +589,7 @@ impl Compiler {
588589
args: &ast::Parameters,
589590
body: &[ast::LocatedStatement],
590591
decorator_list: &[ast::Expression],
592+
_returns: &Option<ast::Expression>, // TODO: use type hint somehow..
591593
) -> Result<(), CompileError> {
592594
// Create bytecode for this function:
593595
// remember to restore self.in_loop to the original after the function is compiled

‎vm/src/stdlib/ast.rs‎

Lines changed: 25 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -82,6 +82,7 @@ fn statement_to_ast(ctx: &PyContext, statement: &ast::LocatedStatement) -> PyObj
8282
args,
8383
body,
8484
decorator_list,
85+
returns,
8586
} => {
8687
let node = create_node(ctx, "FunctionDef");
8788

@@ -96,6 +97,13 @@ fn statement_to_ast(ctx: &PyContext, statement: &ast::LocatedStatement) -> PyObj
9697

9798
let py_decorator_list = expressions_to_ast(ctx, decorator_list);
9899
ctx.set_attr(&node, "decorator_list", py_decorator_list);
100+
101+
let py_returns = if let Some(hint) = returns {
102+
expression_to_ast(ctx, hint)
103+
} else {
104+
ctx.none()
105+
};
106+
ctx.set_attr(&node, "returns", py_returns);
99107
node
100108
}
101109
ast::Statement::Continue => create_node(ctx, "Continue"),
@@ -537,17 +545,28 @@ fn parameters_to_ast(ctx: &PyContext, args: &ast::Parameters) -> PyObjectRef {
537545
ctx.set_attr(
538546
&node,
539547
"args",
540-
ctx.new_list(
541-
args.args
542-
.iter()
543-
.map(|a| ctx.new_str(a.arg.to_string()))
544-
.collect(),
545-
),
548+
ctx.new_list(args.args.iter().map(|a| parameter_to_ast(ctx, a)).collect()),
546549
);
547550

548551
node
549552
}
550553

554+
fn parameter_to_ast(ctx: &PyContext, parameter: &ast::Parameter) -> PyObjectRef {
555+
let node = create_node(ctx, "arg");
556+
557+
let py_arg = ctx.new_str(parameter.arg.to_string());
558+
ctx.set_attr(&node, "arg", py_arg);
559+
560+
let py_annotation = if let Some(annotation) = &parameter.annotation {
561+
expression_to_ast(ctx, annotation)
562+
} else {
563+
ctx.none()
564+
};
565+
ctx.set_attr(&node, "annotation", py_annotation);
566+
567+
node
568+
}
569+
551570
fn comprehension_to_ast(ctx: &PyContext, comprehension: &ast::Comprehension) -> PyObjectRef {
552571
let node = create_node(ctx, "comprehension");
553572

0 commit comments

Comments
 (0)