Skip to content

Commit d278fb8

Browse files
cthulahoopscthulahoops
authored
Merge pull request RustPython#597 from RustPython/type_hints
Add type annotations to parser.
2 parents b7c03b4 + f7a2225 commit d278fb8

File tree

6 files changed

+112
-40
lines changed

6 files changed

+112
-40
lines changed

‎parser/src/ast.rs

Lines changed: 11 additions & 4 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

@@ -271,14 +272,20 @@ impl Expression {
271272
*/
272273
#[derive(Debug, PartialEq, Default)]
273274
pub struct Parameters {
274-
pub args: Vec<String>,
275-
pub kwonlyargs: Vec<String>,
276-
pub vararg: Option<Option<String>>, // Optionally we handle optionally named '*args' or '*'
277-
pub kwarg: Option<Option<String>>,
275+
pub args: Vec<Parameter>,
276+
pub kwonlyargs: Vec<Parameter>,
277+
pub vararg: Option<Option<Parameter>>, // Optionally we handle optionally named '*args' or '*'
278+
pub kwarg: Option<Option<Parameter>>,
278279
pub defaults: Vec<Expression>,
279280
pub kw_defaults: Vec<Option<Expression>>,
280281
}
281282

283+
#[derive(Debug, PartialEq, Default)]
284+
pub struct Parameter {
285+
pub arg: String,
286+
pub annotation: Option<Box<Expression>>,
287+
}
288+
282289
#[derive(Debug, PartialEq)]
283290
pub enum ComprehensionKind {
284291
GeneratorExpression { element: Expression },

‎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: 32 additions & 22 deletions
Original file line numberDiff line numberDiff line change
@@ -446,21 +446,22 @@ 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
},
460461
};
461462

462463
Parameters: ast::Parameters = {
463-
"(" <a: (TypedArgsList)?> ")" => {
464+
"(" <a: (TypedArgsList<TypedParameter>)?> ")" => {
464465
match a {
465466
Some(a) => a,
466467
None => Default::default(),
@@ -470,8 +471,10 @@ Parameters: ast::Parameters = {
470471

471472
// parameters are (String, None), kwargs are (String, Some(Test)) where Test is
472473
// the default
473-
TypedArgsList: ast::Parameters = {
474-
<param1:TypedParameters> <args2:("," ParameterListStarArgs)?> => {
474+
// Note that this is a macro which is used once for function defs, and
475+
// once for lambda defs.
476+
TypedArgsList<ArgType>: ast::Parameters = {
477+
<param1:TypedParameters<ArgType>> <args2:("," ParameterListStarArgs<ArgType>)?> => {
475478
let (names, default_elements) = param1;
476479

477480
// Now gather rest of parameters:
@@ -489,7 +492,7 @@ TypedArgsList: ast::Parameters = {
489492
kw_defaults: kw_defaults,
490493
}
491494
},
492-
<param1:TypedParameters> <kw:("," KwargParameter)> => {
495+
<param1:TypedParameters<ArgType>> <kw:("," KwargParameter<ArgType>)> => {
493496
let (names, default_elements) = param1;
494497

495498
// Now gather rest of parameters:
@@ -507,7 +510,7 @@ TypedArgsList: ast::Parameters = {
507510
kw_defaults: kw_defaults,
508511
}
509512
},
510-
<params:ParameterListStarArgs> => {
513+
<params:ParameterListStarArgs<ArgType>> => {
511514
let (vararg, kwonlyargs, kw_defaults, kwarg) = params;
512515
ast::Parameters {
513516
args: vec![],
@@ -518,7 +521,7 @@ TypedArgsList: ast::Parameters = {
518521
kw_defaults: kw_defaults,
519522
}
520523
},
521-
<kw:KwargParameter> => {
524+
<kw:KwargParameter<ArgType>> => {
522525
ast::Parameters {
523526
args: vec![],
524527
kwonlyargs: vec![],
@@ -532,8 +535,8 @@ TypedArgsList: ast::Parameters = {
532535

533536
// Use inline here to make sure the "," is not creating an ambiguity.
534537
#[inline]
535-
TypedParameters: (Vec<String>, Vec<ast::Expression>) = {
536-
<param1:TypedParameterDef> <param2:("," TypedParameterDef)*> => {
538+
TypedParameters<ArgType>: (Vec<ast::Parameter>, Vec<ast::Expression>) = {
539+
<param1:TypedParameterDef<ArgType>> <param2:("," TypedParameterDef<ArgType>)*> => {
537540
// Combine first parameters:
538541
let mut args = vec![param1];
539542
args.extend(param2.into_iter().map(|x| x.1));
@@ -542,7 +545,6 @@ TypedParameters: (Vec<String>, Vec<ast::Expression>) = {
542545
let mut default_elements = vec![];
543546

544547
for (name, default) in args.into_iter() {
545-
names.push(name.clone());
546548
if let Some(default) = default {
547549
default_elements.push(default);
548550
} else {
@@ -551,28 +553,35 @@ TypedParameters: (Vec<String>, Vec<ast::Expression>) = {
551553
// have defaults
552554
panic!(
553555
"non-default argument follows default argument: {}",
554-
name
556+
&name.arg
555557
);
556558
}
557559
}
560+
names.push(name);
558561
}
559562

560563
(names, default_elements)
561564
}
562565
};
563566

564-
TypedParameterDef: (String, Option<ast::Expression>) = {
565-
<i:TypedParameter> => (i, None),
566-
<i:TypedParameter> "=" <e:Test> => (i, Some(e)),
567+
TypedParameterDef<ArgType>: (ast::Parameter, Option<ast::Expression>) = {
568+
<i:ArgType> => (i, None),
569+
<i:ArgType> "=" <e:Test> => (i, Some(e)),
567570
};
568571

569-
// TODO: add type annotations here:
570-
TypedParameter: String = {
571-
Identifier,
572+
UntypedParameter: ast::Parameter = {
573+
<i:Identifier> => ast::Parameter { arg: i, annotation: None },
572574
};
573575

574-
ParameterListStarArgs: (Option<Option<String>>, Vec<String>, Vec<Option<ast::Expression>>, Option<Option<String>>) = {
575-
"*" <va:Identifier?> <kw:("," TypedParameterDef)*> <kwarg:("," KwargParameter)?> => {
576+
TypedParameter: ast::Parameter = {
577+
<arg:Identifier> <a:(":" Test)?>=> {
578+
let annotation = a.map(|x| Box::new(x.1));
579+
ast::Parameter { arg, annotation }
580+
},
581+
};
582+
583+
ParameterListStarArgs<ArgType>: (Option<Option<ast::Parameter>>, Vec<ast::Parameter>, Vec<Option<ast::Expression>>, Option<Option<ast::Parameter>>) = {
584+
"*" <va:ArgType?> <kw:("," TypedParameterDef<ArgType>)*> <kwarg:("," KwargParameter<ArgType>)?> => {
576585
// Extract keyword arguments:
577586
let mut kwonlyargs = vec![];
578587
let mut kw_defaults = vec![];
@@ -590,8 +599,8 @@ ParameterListStarArgs: (Option<Option<String>>, Vec<String>, Vec<Option<ast::Exp
590599
}
591600
};
592601

593-
KwargParameter: Option<String> = {
594-
"**" <kwarg:Identifier?> => {
602+
KwargParameter<ArgType>: Option<ast::Parameter> = {
603+
"**" <kwarg:ArgType?> => {
595604
kwarg
596605
}
597606
};
@@ -675,7 +684,7 @@ Test: ast::Expression = {
675684
};
676685

677686
LambdaDef: ast::Expression = {
678-
"lambda" <p:TypedArgsList?> ":" <b:Test> =>
687+
"lambda" <p:TypedArgsList<UntypedParameter>?> ":" <b:Test> =>
679688
ast::Expression::Lambda {
680689
args: p.unwrap_or(Default::default()),
681690
body:Box::new(b)
@@ -1088,6 +1097,7 @@ extern {
10881097
"<=" => lexer::Tok::LessEqual,
10891098
">" => lexer::Tok::Greater,
10901099
">=" => lexer::Tok::GreaterEqual,
1100+
"->" => lexer::Tok::Rarrow,
10911101
"and" => lexer::Tok::And,
10921102
"as" => lexer::Tok::As,
10931103
"assert" => lexer::Tok::Assert,

‎tests/snippets/type_hints.py

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,7 @@
1+
2+
# See also: https://github.com/RustPython/RustPython/issues/587
3+
4+
def curry(foo: int) -> float:
5+
return foo * 3.1415926 * 2
6+
7+
assert curry(2) > 10

‎vm/src/compile.rs

Lines changed: 11 additions & 5 deletions
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,
@@ -442,10 +443,14 @@ impl Compiler {
442443

443444
let line_number = self.get_source_line_number();
444445
self.code_object_stack.push(CodeObject::new(
445-
args.args.clone(),
446-
args.vararg.clone(),
447-
args.kwonlyargs.clone(),
448-
args.kwarg.clone(),
446+
args.args.iter().map(|a| a.arg.clone()).collect(),
447+
args.vararg
448+
.as_ref()
449+
.map(|x| x.as_ref().map(|a| a.arg.clone())),
450+
args.kwonlyargs.iter().map(|a| a.arg.clone()).collect(),
451+
args.kwarg
452+
.as_ref()
453+
.map(|x| x.as_ref().map(|a| a.arg.clone())),
449454
self.source_path.clone().unwrap(),
450455
line_number,
451456
name.to_string(),
@@ -584,6 +589,7 @@ impl Compiler {
584589
args: &ast::Parameters,
585590
body: &[ast::LocatedStatement],
586591
decorator_list: &[ast::Expression],
592+
_returns: &Option<ast::Expression>, // TODO: use type hint somehow..
587593
) -> Result<(), CompileError> {
588594
// Create bytecode for this function:
589595
// 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"),
@@ -538,17 +546,28 @@ fn parameters_to_ast(ctx: &PyContext, args: &ast::Parameters) -> PyObjectRef {
538546
ctx.set_attr(
539547
&node,
540548
"args",
541-
ctx.new_list(
542-
args.args
543-
.iter()
544-
.map(|a| ctx.new_str(a.to_string()))
545-
.collect(),
546-
),
549+
ctx.new_list(args.args.iter().map(|a| parameter_to_ast(ctx, a)).collect()),
547550
);
548551

549552
node
550553
}
551554

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

0 commit comments

Comments
 (0)