bpo-40334: Support type comments

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") }

Lib/test/test_type_comments.py

@@ -219,7 +219,6 @@ def favk(
 """
 
 
-@support.skip_if_new_parser("Pegen does not support type comments yet")
 class TypeCommentTests(unittest.TestCase):
 
     lowest = 4  # Lowest minor version supported
@@ -253,6 +252,7 @@ def test_funcdef(self):
         self.assertEqual(tree.body[0].type_comment, None)
         self.assertEqual(tree.body[1].type_comment, None)
 
+    @support.skip_if_new_parser("Pegen does not support feature_version yet")
     def test_asyncdef(self):
         for tree in self.parse_all(asyncdef, minver=5):
             self.assertEqual(tree.body[0].type_comment, "() -> int")
@@ -261,22 +261,27 @@ def test_asyncdef(self):
         self.assertEqual(tree.body[0].type_comment, None)
         self.assertEqual(tree.body[1].type_comment, None)
 
+    @support.skip_if_new_parser("Pegen does not support feature_version yet")
     def test_asyncvar(self):
         for tree in self.parse_all(asyncvar, maxver=6):
             pass
 
+    @support.skip_if_new_parser("Pegen does not support feature_version yet")
     def test_asynccomp(self):
         for tree in self.parse_all(asynccomp, minver=6):
             pass
 
+    @support.skip_if_new_parser("Pegen does not support feature_version yet")
     def test_matmul(self):
         for tree in self.parse_all(matmul, minver=5):
             pass
 
+    @support.skip_if_new_parser("Pegen does not support feature_version yet")
     def test_fstring(self):
         for tree in self.parse_all(fstring, minver=6):
             pass
 
+    @support.skip_if_new_parser("Pegen does not support feature_version yet")
     def test_underscorednumber(self):
         for tree in self.parse_all(underscorednumber, minver=6):
             pass