Add support for parsing type-comment in func and async func.

yashsain · Jul 20, 2022 · bcb8639 · bcb8639
1 parent e1d0491
commit bcb8639
Show file tree

Hide file tree

Showing 2 changed files with 39 additions and 0 deletions.
diff --git a/src/lpython/parser/parser.yy b/src/lpython/parser/parser.yy
@@ -614,6 +614,11 @@ function_def
         $$ = FUNCTION_01($1, $3, $5, $9, @$); }
     | decorators_opt KW_DEF id "(" parameter_list_opt ")" "->" expr ":"
         sep statements { $$ = FUNCTION_02($1, $3, $5, $8, $11, @$); }
+    | decorators_opt KW_DEF id "(" parameter_list_opt ")" ":" TK_TYPE_COMMENT TK_NEWLINE
+        statements { $$ = FUNCTION_03($1, $3, $5, $10, $8, @$); }
+    | decorators_opt KW_DEF id "(" parameter_list_opt ")" "->" expr ":"
+        TK_TYPE_COMMENT TK_NEWLINE statements {
+            $$ = FUNCTION_04($1, $3, $5, $8, $12, $10, @$); }
     ;
 
 class_def
@@ -631,6 +636,16 @@ async_func_def
         statements { $$ = ASYNC_FUNCTION_03($3, $5, $9, @$); }
     | KW_ASYNC KW_DEF id "(" parameter_list_opt ")" "->" expr ":"
         sep statements { $$ = ASYNC_FUNCTION_04($3, $5, $8, $11, @$); }
+    | decorators KW_ASYNC KW_DEF id "(" parameter_list_opt ")" ":" TK_TYPE_COMMENT
+        TK_NEWLINE statements { $$ = ASYNC_FUNCTION_05($1, $4, $6, $11, $9, @$); }
+    | decorators KW_ASYNC KW_DEF id "(" parameter_list_opt ")" "->" expr ":"
+        TK_TYPE_COMMENT TK_NEWLINE statements {
+        $$ = ASYNC_FUNCTION_06($1, $4, $6, $9, $13, $11, @$); }
+    | KW_ASYNC KW_DEF id "(" parameter_list_opt ")" ":" TK_TYPE_COMMENT TK_NEWLINE
+        statements { $$ = ASYNC_FUNCTION_07($3, $5, $10, $8, @$); }
+    | KW_ASYNC KW_DEF id "(" parameter_list_opt ")" "->" expr ":"
+        TK_TYPE_COMMENT TK_NEWLINE statements {
+        $$ = ASYNC_FUNCTION_08($3, $5, $8, $12, $10, @$); }
     ;
 
 async_for_stmt

diff --git a/src/lpython/parser/semantics.h b/src/lpython/parser/semantics.h
@@ -425,6 +425,14 @@ static inline Args *FUNC_ARGS(Allocator &al, Location &l,
         make_FunctionDef_t(p.m_a, l, name2char(id), args->arguments, \
         STMTS(stmts), stmts.size(), EXPRS(decorator), decorator.size(), \
         EXPR(return), nullptr)
+#define FUNCTION_03(decorator, id, args, stmts, type_comment, l) \
+        make_FunctionDef_t(p.m_a, l, name2char(id), args->arguments, \
+        STMTS(stmts), stmts.size(), EXPRS(decorator), decorator.size(), \
+        nullptr, extract_type_comment(p.m_a, type_comment))
+#define FUNCTION_04(decorator, id, args, return, stmts, type_comment, l) \
+        make_FunctionDef_t(p.m_a, l, name2char(id), args->arguments, \
+        STMTS(stmts), stmts.size(), EXPRS(decorator), decorator.size(), \
+        EXPR(return), extract_type_comment(p.m_a, type_comment))
 
 #define CLASS_01(decorator, id, stmts, l) make_ClassDef_t(p.m_a, l, \
         name2char(id), nullptr, 0, nullptr, 0, \
@@ -449,6 +457,22 @@ static inline Args *FUNC_ARGS(Allocator &al, Location &l,
 #define ASYNC_FUNCTION_04(id, args, return, stmts, l) \
         make_AsyncFunctionDef_t(p.m_a, l, name2char(id), args->arguments, \
         STMTS(stmts), stmts.size(), nullptr, 0, EXPR(return), nullptr)
+#define ASYNC_FUNCTION_05(decorator, id, args, stmts, type_comment, l) \
+        make_AsyncFunctionDef_t(p.m_a, l, name2char(id), args->arguments, \
+        STMTS(stmts), stmts.size(), EXPRS(decorator), decorator.size(), \
+        nullptr, extract_type_comment(p.m_a, type_comment))
+#define ASYNC_FUNCTION_06(decorator, id, args, return, stmts, type_comment, l) \
+        make_AsyncFunctionDef_t(p.m_a, l, name2char(id), args->arguments, \
+        STMTS(stmts), stmts.size(), EXPRS(decorator), decorator.size(), \
+        EXPR(return), extract_type_comment(p.m_a, type_comment))
+#define ASYNC_FUNCTION_07(id, args, stmts, type_comment, l) \
+        make_AsyncFunctionDef_t(p.m_a, l, name2char(id), args->arguments, \
+        STMTS(stmts), stmts.size(), nullptr, 0, nullptr,\
+        extract_type_comment(p.m_a, type_comment))
+#define ASYNC_FUNCTION_08(id, args, return, stmts, type_comment, l) \
+        make_AsyncFunctionDef_t(p.m_a, l, name2char(id), args->arguments, \
+        STMTS(stmts), stmts.size(), nullptr, 0, EXPR(return),\
+        extract_type_comment(p.m_a, type_comment))
 
 #define ASYNC_FOR_01(target, iter, stmts, l) make_AsyncFor_t(p.m_a, l, \
         EXPR(SET_EXPR_CTX_01(SET_STORE_01(target), Store)), EXPR(iter), \