Merge pull request lcompilers#933 from redbopo/issue-920

Fix issue-920

src/lpython/parser/parser.yy

@@ -486,6 +486,7 @@ dot_list
 
 import_statement
     : KW_IMPORT module_item_list { $$ = IMPORT_01($2, @$); }
+    | KW_IMPORT module_item_list TK_TYPE_COMMENT { $$ = IMPORT_01($2, @$); }
     | KW_FROM module KW_IMPORT module_item_list {
         $$ = IMPORT_02($2, $4, @$); }
     | KW_FROM module KW_IMPORT "(" module_item_list comma_opt ")" {
@@ -667,6 +668,12 @@ function_def
     | decorators_opt KW_DEF id "(" parameter_list_opt ")" "->" expr ":"
         TK_TYPE_COMMENT TK_NEWLINE statements {
             $$ = FUNCTION_04($1, $3, $5, $8, $12, $10, @$); }
+    | decorators_opt KW_DEF id "(" parameter_list_opt ")" ":"
+        TK_NEWLINE TK_TYPE_COMMENT TK_NEWLINE statements {
+        $$ = FUNCTION_03($1, $3, $5, $11, $9, @$); }
+    | decorators_opt KW_DEF id "(" parameter_list_opt ")" "->" expr ":"
+        TK_NEWLINE TK_TYPE_COMMENT TK_NEWLINE statements {
+            $$ = FUNCTION_04($1, $3, $5, $8, $13, $11, @$); }
     ;
 
 class_def

tests/parser/type_comment1.py

@@ -0,0 +1,5 @@
+import pytest  # type: ignore
+
+def ndarray_func(x):
+    # type: (np.ndarray) -> np.ndarray
+    return x

tests/reference/ast_new-type_comment1-710ea6c.json

@@ -0,0 +1,13 @@
+{
+    "basename": "ast_new-type_comment1-710ea6c",
+    "cmd": "lpython --show-ast --new-parser --no-color {infile} -o {outfile}",
+    "infile": "tests/parser/type_comment1.py",
+    "infile_hash": "cdc0bdd0faec912e7f3db4e71c4a14aa281828e8b68faac91f5a78ee",
+    "outfile": null,
+    "outfile_hash": null,
+    "stdout": "ast_new-type_comment1-710ea6c.stdout",
+    "stdout_hash": "98ccc7f80021238fff1a1a1395e64b5d2097f025e60e47821ec76fdb",
+    "stderr": null,
+    "stderr_hash": null,
+    "returncode": 0
+}

tests/reference/ast_new-type_comment1-710ea6c.stdout

@@ -0,0 +1 @@
+(Module [(Import [(pytest ())]) (FunctionDef ndarray_func ([] [(x () ())] [] [] [] [] []) [(Return (Name x Load))] [] () "(np.ndarray) -> np.ndarray")] [])

tests/tests.toml

@@ -524,6 +524,10 @@ ast_new = true
 filename = "parser/tuple1.py"
 ast_new = true
 
+[[test]]
+filename = "parser/type_comment1.py"
+ast_new = true
+
 # tests/errors
 
 [[test]]