Limit match range to end of last statement

Author: charliermarsh

parser/python.lalrpop

@@ -359,7 +359,15 @@ CompoundStatement: ast::Stmt = {
 };
 
 MatchStatement: ast::Stmt = {
-    <location:@L> "match" <subject:TestOrStarNamedExpr> ":" "\n" Indent <cases:MatchCase+> Dedent <end_location:@R> => {
+    <location:@L> "match" <subject:TestOrStarNamedExpr> ":" "\n" Indent <cases:MatchCase+> Dedent => {
+       let end_location = cases
+            .last()
+            .unwrap()
+            .body
+            .last()
+            .unwrap()
+            .end_location
+            .unwrap();
         ast::Stmt {
             location,
             end_location: Some(end_location),
@@ -370,7 +378,15 @@ MatchStatement: ast::Stmt = {
             }
         }
     },
-    <location:@L> "match" <subject:TestOrStarNamedExpr> "," ":" "\n" Indent <cases:MatchCase+> Dedent <end_location:@R> => {
+    <location:@L> "match" <subject:TestOrStarNamedExpr> "," ":" "\n" Indent <cases:MatchCase+> Dedent => {
+       let end_location = cases
+            .last()
+            .unwrap()
+            .body
+            .last()
+            .unwrap()
+            .end_location
+            .unwrap();
         ast::Stmt {
             location,
             end_location: Some(end_location),
@@ -381,7 +397,15 @@ MatchStatement: ast::Stmt = {
             }
         }
     },
-    <location:@L> "match" <subject:TestOrStarNamedExpr> "," <subjects:OneOrMore<TestOrStarNamedExpr>> ","? ":" "\n" Indent <cases:MatchCase+> Dedent <end_location:@R> => {
+    <location:@L> "match" <subject:TestOrStarNamedExpr> "," <subjects:OneOrMore<TestOrStarNamedExpr>> ","? ":" "\n" Indent <cases:MatchCase+> Dedent => {
+        let end_location = cases
+            .last()
+            .unwrap()
+            .body
+            .last()
+            .unwrap()
+            .end_location
+            .unwrap();
         let mut subjects = subjects;
         subjects.insert(0, subject);
         ast::Stmt {

parser/src/snapshots/rustpython_parser__parser__tests__match.snap

@@ -10,8 +10,8 @@ expression: parse_ast
         },
         end_location: Some(
             Location {
-                row: 7,
-                column: 0,
+                row: 6,
+                column: 19,
             },
         ),
         custom: (),
@@ -178,8 +178,8 @@ expression: parse_ast
         },
         end_location: Some(
             Location {
-                row: 12,
-                column: 0,
+                row: 11,
+                column: 20,
             },
         ),
         custom: (),
@@ -453,8 +453,8 @@ expression: parse_ast
         },
         end_location: Some(
             Location {
-                row: 15,
-                column: 0,
+                row: 14,
+                column: 13,
             },
         ),
         custom: (),
@@ -631,8 +631,8 @@ expression: parse_ast
         },
         end_location: Some(
             Location {
-                row: 18,
-                column: 0,
+                row: 17,
+                column: 13,
             },
         ),
         custom: (),
@@ -809,8 +809,8 @@ expression: parse_ast
         },
         end_location: Some(
             Location {
-                row: 21,
-                column: 0,
+                row: 20,
+                column: 13,
             },
         ),
         custom: (),

parser/src/snapshots/rustpython_parser__parser__tests__match_as_identifier.snap

@@ -1538,8 +1538,8 @@ expression: parse_ast
         },
         end_location: Some(
             Location {
-                row: 21,
-                column: 0,
+                row: 20,
+                column: 12,
             },
         ),
         custom: (),