Merge pull request #993 from akshanshbhatt/pr_compr_subr

[Parser] Recognise `ListComprehension` Subscript

Author: akshanshbhatt

src/lpython/parser/parser.yy

@@ -250,7 +250,7 @@ void yyerror(YYLTYPE *yyloc, LFortran::Parser &p, const std::string &msg)
 %type <string> type_comment
 %type <ast> string
 %type <ast> ternary_if_statement
-/* %type <ast> list_comprehension */
+%type <ast> comprehension
 %type <vec_ast> id_list
 %type <vec_ast> id_list_opt
 %type <ast> id_item
@@ -908,6 +908,7 @@ subscript
     | "{" dict_list comma_opt "}" "[" slice_item "]" {
         $$ = SUBSCRIPT_01(DICT_02($2, @$), $6, @$); }
     | subscript "[" slice_item "]" { $$ = SUBSCRIPT_01($1, $3, @$); }
+    | comprehension "[" slice_item "]" { $$ = SUBSCRIPT_01($1, $3, @$); }
     ;
 
 string
@@ -968,6 +969,13 @@ lambda_expression
         $$ = LAMBDA_01($2, $4, @$); }
     ;
 
+comprehension
+    : "[" expr comp_for_items "]" { $$ = LIST_COMP_1($2, $3, @$); }
+    | "{" expr comp_for_items "}" { $$ = SET_COMP_1($2, $3, @$); }
+    | "{" expr ":" expr comp_for_items "}" { $$ = DICT_COMP_1($2, $4, $5, @$); }
+    | "(" expr comp_for_items ")" { $$ = COMP_EXPR_1($2, $3, @$); }
+    ;
+
 expr
     : id { $$ = $1; }
     | TK_INTEGER { $$ = INTEGER($1, @$); }
@@ -1028,12 +1036,7 @@ expr
     | expr "or" expr { $$ = BOOLOP($1, Or, $3, @$); }
     | "not" expr { $$ = UNARY($2, Not, @$); }
 
-    // Comprehension
-    | "[" expr comp_for_items "]" { $$ = LIST_COMP_1($2, $3, @$); }
-    | "{" expr comp_for_items "}" { $$ = SET_COMP_1($2, $3, @$); }
-    | "{" expr ":" expr comp_for_items "}" { $$ = DICT_COMP_1($2, $4, $5, @$); }
-    | "(" expr comp_for_items ")" { $$ = COMP_EXPR_1($2, $3, @$); }
-
+    | comprehension { $$ = $1; }
     | ternary_if_statement { $$ = $1; }
     | lambda_expression { $$ = $1; }
     ;

tests/parser/statements2.py

@@ -14,3 +14,7 @@
 
 if args[i][:1] in ['', '.']:
     pass
+
+x = [[-c % self for c in reversed(T.rep.rep)][:-1]]
+
+[x for x in G if self.ring.is_unit(x[0])][0]

tests/reference/ast_new-statements2-c4cdc5f.json

@@ -2,11 +2,11 @@
     "basename": "ast_new-statements2-c4cdc5f",
     "cmd": "lpython --show-ast --new-parser --no-color {infile} -o {outfile}",
     "infile": "tests/parser/statements2.py",
-    "infile_hash": "582c1146456415a483654db39eb660b14593a6e72f353a61ed716b4b",
+    "infile_hash": "e16da92e6b14ae0db259cf323ccb98390c5fcee8f199914e3f1a9525",
     "outfile": null,
     "outfile_hash": null,
     "stdout": "ast_new-statements2-c4cdc5f.stdout",
-    "stdout_hash": "d9c72edbe478ee87b62391369aba8651ca66af28000332f6a38eac37",
+    "stdout_hash": "6c1548bf54cb810ff6bc25816b0a63484aa60192e6fe246df412cf50",
     "stderr": null,
     "stderr_hash": null,
     "returncode": 0

tests/reference/ast_new-statements2-c4cdc5f.stdout

@@ -1 +1 @@
-(Module [(Expr (Call (Call (Call (Name getattr Load) [(Name x Load) (Name y Load)] []) [(Name a Load) (Name b Load)] []) [(ConstantInt 5 ())] [(a (Name a Load)) (b (ConstantInt 0 ()))])) (Expr (Call (Call (Name func Load) [(Name a Load) (Name _op Load)] []) [(Name s Load)] [])) (Expr (Call (Call (Name test Load) [] []) [] [])) (Assign [(Name a Store)] (Subscript (List [(ConstantInt 1 ()) (ConstantInt 2 ())] Load) (ConstantInt 0 ()) Load) ()) (Assign [(Tuple [(Name a Store) (Name b Store)] Store)] (Subscript (BinOp (List [(ConstantInt 1 ()) (ConstantInt 2 ())] Load) Add (List [(ConstantInt 0 ()) (ConstantInt 0 ())] Load)) (Slice () (ConstantInt 2 ()) ()) Load) ()) (Assign [(Tuple [(Name a Store) (Name b Store)] Store)] (Subscript (BinOp (Call (Attribute (Name c Load) d Load) [] []) Add (List [(ConstantInt 1 ()) (ConstantInt 0 ())] Load)) (Slice () () ()) Load) ()) (Expr (Subscript (Dict [(ConstantStr "a" ()) (ConstantStr "b" ())] [(Name a Load) (Name b Load)]) (Name val Load) Load)) (If (Compare (Subscript (Subscript (Name args Load) (Name i Load) Load) (Slice () (ConstantInt 1 ()) ()) Load) In [(List [(ConstantStr "" ()) (ConstantStr "." ())] Load)]) [(Pass)] [])] [])
+(Module [(Expr (Call (Call (Call (Name getattr Load) [(Name x Load) (Name y Load)] []) [(Name a Load) (Name b Load)] []) [(ConstantInt 5 ())] [(a (Name a Load)) (b (ConstantInt 0 ()))])) (Expr (Call (Call (Name func Load) [(Name a Load) (Name _op Load)] []) [(Name s Load)] [])) (Expr (Call (Call (Name test Load) [] []) [] [])) (Assign [(Name a Store)] (Subscript (List [(ConstantInt 1 ()) (ConstantInt 2 ())] Load) (ConstantInt 0 ()) Load) ()) (Assign [(Tuple [(Name a Store) (Name b Store)] Store)] (Subscript (BinOp (List [(ConstantInt 1 ()) (ConstantInt 2 ())] Load) Add (List [(ConstantInt 0 ()) (ConstantInt 0 ())] Load)) (Slice () (ConstantInt 2 ()) ()) Load) ()) (Assign [(Tuple [(Name a Store) (Name b Store)] Store)] (Subscript (BinOp (Call (Attribute (Name c Load) d Load) [] []) Add (List [(ConstantInt 1 ()) (ConstantInt 0 ())] Load)) (Slice () () ()) Load) ()) (Expr (Subscript (Dict [(ConstantStr "a" ()) (ConstantStr "b" ())] [(Name a Load) (Name b Load)]) (Name val Load) Load)) (If (Compare (Subscript (Subscript (Name args Load) (Name i Load) Load) (Slice () (ConstantInt 1 ()) ()) Load) In [(List [(ConstantStr "" ()) (ConstantStr "." ())] Load)]) [(Pass)] []) (Assign [(Name x Store)] (List [(Subscript (ListComp (BinOp (UnaryOp USub (Name c Load)) Mod (Name self Load)) [((Name c Store) (Call (Name reversed Load) [(Attribute (Attribute (Name T Load) rep Load) rep Load)] []) [] 0)]) (Slice () (UnaryOp USub (ConstantInt 1 ())) ()) Load)] Load) ()) (Expr (Subscript (ListComp (Name x Load) [((Name x Store) (Name G Load) [(Call (Attribute (Attribute (Name self Load) ring Load) is_unit Load) [(Subscript (Name x Load) (ConstantInt 0 ()) Load)] [])] 0)]) (ConstantInt 0 ()) Load))] [])