bpo-40334: Support type comments (GH-19780)

This implements full support for # type: <type> comments, # type: ignore <stuff> comments, and the func_type parsing mode for ast.parse() and compile().

Closes https://github.com/we-like-parsers/cpython/issues/95.

(For now, you need to use the master branch of mypy, since another issue unique to 3.9 had to be fixed there, and there's no mypy release yet.)

The only thing missing is `feature_version=N`, which is being tracked in https://github.com/we-like-parsers/cpython/issues/124.

1 files changed, 77 insertions, 32 deletions

diff --git a/Grammar/python.gram b/Grammar/python.gram
index 0ff2dcc..e1164d9 100644
--- a/Grammar/python.gram
+++ b/Grammar/python.gram
@@ -17,6 +17,8 @@ _PyPegen_parse(Parser *p)
         result = interactive_rule(p);
     } else if (p->start_rule == Py_eval_input) {
         result = eval_rule(p);
+    } else if (p->start_rule == Py_func_type_input) {
+        result = func_type_rule(p);
     } else if (p->start_rule == Py_fstring_input) {
         result = fstring_rule(p);
     }
@@ -26,11 +28,20 @@ _PyPegen_parse(Parser *p)
 
 // The end
 '''
-file[mod_ty]: a=[statements] ENDMARKER { Module(a, NULL, p->arena) }
+file[mod_ty]: a=[statements] ENDMARKER { _PyPegen_make_module(p, a) }
 interactive[mod_ty]: a=statement_newline { Interactive(a, p->arena) }
 eval[mod_ty]: a=expressions NEWLINE* ENDMARKER { Expression(a, p->arena) }
+func_type[mod_ty]: '(' a=[type_expressions] ')' '->' b=expression NEWLINE* ENDMARKER { FunctionType(a, b, p->arena) }
 fstring[expr_ty]: star_expressions
 
+# type_expressions allow */** but ignore them
+type_expressions[asdl_seq*]:
+    | a=','.expression+ ',' '*' b=expression ',' '**' c=expression {
+        _PyPegen_seq_append_to_end(p, CHECK(_PyPegen_seq_append_to_end(p, a, b)), c) }
+    | a=','.expression+ ',' '*' b=expression { _PyPegen_seq_append_to_end(p, a, b) }
+    | a=','.expression+ ',' '**' b=expression { _PyPegen_seq_append_to_end(p, a, b) }
+    | ','.expression+
+
 statements[asdl_seq*]: a=statement+ { _PyPegen_seq_flatten(p, a) }
 statement[asdl_seq*]: a=compound_stmt { _PyPegen_singleton_seq(p, a) } | simple_stmt
 statement_newline[asdl_seq*]:
@@ -73,8 +84,8 @@ assignment:
     | a=('(' b=inside_paren_ann_assign_target ')' { b }
          | ann_assign_subscript_attribute_target) ':' b=expression c=['=' d=annotated_rhs { d }] {
         _Py_AnnAssign(a, b, c, 0, EXTRA)}
-    | a=(z=star_targets '=' { z })+ b=(yield_expr | star_expressions) {
-         _Py_Assign(a, b, NULL, EXTRA) }
+    | a=(z=star_targets '=' { z })+ b=(yield_expr | star_expressions) tc=[TYPE_COMMENT] {
+         _Py_Assign(a, b, NEW_TYPE_COMMENT(p, tc), EXTRA) }
     | a=target b=augassign c=(yield_expr | star_expressions) {
          _Py_AugAssign(a, b->kind, c, EXTRA) }
     | invalid_assignment
@@ -145,14 +156,14 @@ while_stmt[stmt_ty]:
     | 'while' a=named_expression ':' b=block c=[else_block] { _Py_While(a, b, c, EXTRA) }
 
 for_stmt[stmt_ty]:
-    | is_async=[ASYNC] 'for' t=star_targets 'in' ex=star_expressions ':' b=block el=[else_block] {
-        (is_async ? _Py_AsyncFor : _Py_For)(t, ex, b, el, NULL, EXTRA) }
+    | is_async=[ASYNC] 'for' t=star_targets 'in' ex=star_expressions ':' tc=[TYPE_COMMENT] b=block el=[else_block] {
+        (is_async ? _Py_AsyncFor : _Py_For)(t, ex, b, el, NEW_TYPE_COMMENT(p, tc), EXTRA) }
 
 with_stmt[stmt_ty]:
     | is_async=[ASYNC] 'with' '(' a=','.with_item+ ')' ':' b=block {
         (is_async ? _Py_AsyncWith : _Py_With)(a, b, NULL, EXTRA) }
-    | is_async=[ASYNC] 'with' a=','.with_item+ ':' b=block {
-        (is_async ? _Py_AsyncWith : _Py_With)(a, b, NULL, EXTRA) }
+    | is_async=[ASYNC] 'with' a=','.with_item+ ':' tc=[TYPE_COMMENT] b=block {
+        (is_async ? _Py_AsyncWith : _Py_With)(a, b, NEW_TYPE_COMMENT(p, tc), EXTRA) }
 with_item[withitem_ty]:
     | e=expression o=['as' t=target { t }] { _Py_withitem(e, o, p->arena) }
 
@@ -177,43 +188,74 @@ function_def[stmt_ty]:
     | function_def_raw
 
 function_def_raw[stmt_ty]:
-    | is_async=[ASYNC] 'def' n=NAME '(' params=[params] ')' a=['->' z=annotation { z }] ':' b=block {
+    | is_async=[ASYNC] 'def' n=NAME '(' params=[params] ')' a=['->' z=expression { z }] ':' tc=[func_type_comment] b=block {
         (is_async ? _Py_AsyncFunctionDef : _Py_FunctionDef)(n->v.Name.id,
                              (params) ? params : CHECK(_PyPegen_empty_arguments(p)),
-                             b, NULL, a, NULL, EXTRA) }
+                             b, NULL, a, NEW_TYPE_COMMENT(p, tc), EXTRA) }
+func_type_comment[PyObject*]:
+    | NEWLINE t=TYPE_COMMENT &(NEWLINE INDENT) { t }  # Must be followed by indented block
+    | invalid_double_type_comments
+    | TYPE_COMMENT
 
 params[arguments_ty]:
     | invalid_parameters
     | parameters
+
 parameters[arguments_ty]:
-    | a=slash_without_default b=[',' x=plain_names { x }] c=[',' y=names_with_default { y }] d=[',' z=[star_etc] { z }] {
+    | a=slash_no_default b=param_no_default* c=param_with_default* d=[star_etc] {
         _PyPegen_make_arguments(p, a, NULL, b, c, d) }
-    | a=slash_with_default b=[',' y=names_with_default { y }] c=[',' z=[star_etc] { z }] {
+    | a=slash_with_default b=param_with_default* c=[star_etc] {
         _PyPegen_make_arguments(p, NULL, a, NULL, b, c) }
-    | a=plain_names b=[',' y=names_with_default { y }] c=[',' z=[star_etc] { z }] {
+    | a=param_no_default+ b=param_with_default* c=[star_etc] {
         _PyPegen_make_arguments(p, NULL, NULL, a, b, c) }
-    | a=names_with_default b=[',' z=[star_etc] { z }] { _PyPegen_make_arguments(p, NULL, NULL, NULL, a, b)}
+    | a=param_with_default+ b=[star_etc] { _PyPegen_make_arguments(p, NULL, NULL, NULL, a, b)}
     | a=star_etc { _PyPegen_make_arguments(p, NULL, NULL, NULL, NULL, a) }
-slash_without_default[asdl_seq*]: a=plain_names ',' '/' { a }
-slash_with_default[SlashWithDefault*]: a=[n=plain_names ',' { n }] b=names_with_default ',' '/' {
-    _PyPegen_slash_with_default(p, a, b) }
+
+# Some duplication here because we can't write (',' | &')'),
+# which is because we don't support empty alternatives (yet).
+#
+slash_no_default[asdl_seq*]:
+    | a=param_no_default+ '/' ',' { a }
+    | a=param_no_default+ '/' &')' { a }
+slash_with_default[SlashWithDefault*]:
+    | a=param_no_default* b=param_with_default+ '/' ',' { _PyPegen_slash_with_default(p, a, b) }
+    | a=param_no_default* b=param_with_default+ '/' &')' { _PyPegen_slash_with_default(p, a, b) }
+
 star_etc[StarEtc*]:
-    | '*' a=plain_name b=name_with_optional_default* c=[',' d=kwds { d }] [','] {
+    | '*' a=param_no_default b=param_maybe_default* c=[kwds] {
         _PyPegen_star_etc(p, a, b, c) }
-    | '*' b=name_with_optional_default+ c=[',' d=kwds { d }] [','] {
+    | '*' ',' b=param_maybe_default+ c=[kwds] {
         _PyPegen_star_etc(p, NULL, b, c) }
-    | a=kwds [','] { _PyPegen_star_etc(p, NULL, NULL, a) }
-name_with_optional_default[NameDefaultPair*]:
-    | ',' a=plain_name b=['=' e=expression { e }] { _PyPegen_name_default_pair(p, a, b) }
-names_with_default[asdl_seq*]: a=','.name_with_default+ { a }
-name_with_default[NameDefaultPair*]:
-    | n=plain_name '=' e=expression { _PyPegen_name_default_pair(p, n, e) }
-plain_names[asdl_seq*] (memo): a=','.(plain_name !'=')+ { a }
-plain_name[arg_ty]:
-    | a=NAME b=[':' z=annotation { z }] { _Py_arg(a->v.Name.id, b, NULL, EXTRA) }
+    | a=kwds { _PyPegen_star_etc(p, NULL, NULL, a) }
+
 kwds[arg_ty]:
-    | '**' a=plain_name { a }
-annotation[expr_ty]: expression
+    | '**' a=param_no_default { a }
+
+# One parameter.  This *includes* a following comma and type comment.
+#
+# There are three styles:
+# - No default
+# - With default
+# - Maybe with default
+#
+# There are two alternative forms of each, to deal with type comments:
+# - Ends in a comma followed by an optional type comment
+# - No comma, optional type comment, must be followed by close paren
+# The latter form is for a final parameter without trailing comma.
+#
+param_no_default[arg_ty]:
+    | a=param ',' tc=TYPE_COMMENT? { _PyPegen_add_type_comment_to_arg(p, a, tc) }
+    | a=param tc=TYPE_COMMENT? &')' { _PyPegen_add_type_comment_to_arg(p, a, tc) }
+param_with_default[NameDefaultPair*]:
+    | a=param c=default ',' tc=TYPE_COMMENT? { _PyPegen_name_default_pair(p, a, c, tc) }
+    | a=param c=default tc=TYPE_COMMENT? &')' { _PyPegen_name_default_pair(p, a, c, tc) }
+param_maybe_default[NameDefaultPair*]:
+    | a=param c=default? ',' tc=TYPE_COMMENT? { _PyPegen_name_default_pair(p, a, c, tc) }
+    | a=param c=default? tc=TYPE_COMMENT? &')' { _PyPegen_name_default_pair(p, a, c, tc) }
+param[arg_ty]: a=NAME b=annotation? { _Py_arg(a->v.Name.id, b, NULL, EXTRA) }
+
+annotation[expr_ty]: ':' a=expression { a }
+default[expr_ty]: '=' a=expression { a }
 
 decorators[asdl_seq*]: a=('@' f=named_expression NEWLINE { f })+ { a }
 
@@ -284,10 +326,10 @@ lambda_star_etc[StarEtc*]:
         _PyPegen_star_etc(p, NULL, b, c) }
     | a=lambda_kwds [','] { _PyPegen_star_etc(p, NULL, NULL, a) }
 lambda_name_with_optional_default[NameDefaultPair*]:
-    | ',' a=lambda_plain_name b=['=' e=expression { e }] { _PyPegen_name_default_pair(p, a, b) }
+    | ',' a=lambda_plain_name b=['=' e=expression { e }] { _PyPegen_name_default_pair(p, a, b, NULL) }
 lambda_names_with_default[asdl_seq*]: a=','.lambda_name_with_default+ { a }
 lambda_name_with_default[NameDefaultPair*]:
-    | n=lambda_plain_name '=' e=expression { _PyPegen_name_default_pair(p, n, e) }
+    | n=lambda_plain_name '=' e=expression { _PyPegen_name_default_pair(p, n, e, NULL) }
 lambda_plain_names[asdl_seq*]: a=','.(lambda_plain_name !'=')+ { a }
 lambda_plain_name[arg_ty]: a=NAME { _Py_arg(a->v.Name.id, NULL, NULL, EXTRA) }
 lambda_kwds[arg_ty]: '**' a=lambda_plain_name { a }
@@ -552,5 +594,8 @@ invalid_comprehension:
     | ('[' | '(' | '{') '*' expression for_if_clauses {
         RAISE_SYNTAX_ERROR("iterable unpacking cannot be used in comprehension") }
 invalid_parameters:
-    | [plain_names ','] (slash_with_default | names_with_default) ',' plain_names {
+    | param_no_default* (slash_with_default | param_with_default+) param_no_default {
         RAISE_SYNTAX_ERROR("non-default argument follows default argument") }
+invalid_double_type_comments:
+    | TYPE_COMMENT NEWLINE TYPE_COMMENT NEWLINE INDENT {
+        RAISE_SYNTAX_ERROR("Cannot have two type comments on def") }