Merge pull request #4327 from harupy/fix-end-location-body

Fix end location of compound statements
This commit is contained in:
Jim Fasarakis-Hilliard 2022-12-13 12:30:21 +02:00 committed by GitHub
commit ca1fe40701
5 changed files with 94 additions and 65 deletions

View file

@ -360,16 +360,21 @@ CompoundStatement: ast::Stmt = {
}; };
IfStatement: ast::Stmt = { IfStatement: ast::Stmt = {
<location:@L> "if" <test:NamedExpressionTest> ":" <body:Suite> <s2:(@L "elif" NamedExpressionTest ":" Suite @R)*> <s3:("else" ":" Suite)?> <end_location:@R> => { <location:@L> "if" <test:NamedExpressionTest> ":" <body:Suite> <s2:(@L "elif" NamedExpressionTest ":" Suite)*> <s3:("else" ":" Suite)?> => {
// Determine last else: // Determine last else:
let mut last = s3.map(|s| s.2).unwrap_or_default(); let mut last = s3.map(|s| s.2).unwrap_or_default();
let end_location = last
.last()
.or_else(|| s2.last().and_then(|last| last.4.last()))
.or_else(|| body.last())
.unwrap()
.end_location;
// handle elif: // handle elif:
for i in s2.into_iter().rev() { for i in s2.into_iter().rev() {
let x = ast::Stmt { let x = ast::Stmt {
custom: (), custom: (),
location: i.0, location: i.0,
end_location: Some(i.5), end_location: i.4.last().unwrap().end_location,
node: ast::StmtKind::If { test: Box::new(i.2), body: i.4, orelse: last }, node: ast::StmtKind::If { test: Box::new(i.2), body: i.4, orelse: last },
}; };
last = vec![x]; last = vec![x];
@ -378,19 +383,24 @@ IfStatement: ast::Stmt = {
ast::Stmt { ast::Stmt {
custom: (), custom: (),
location, location,
end_location: Some(end_location), end_location,
node: ast::StmtKind::If { test: Box::new(test), body, orelse: last } node: ast::StmtKind::If { test: Box::new(test), body, orelse: last }
} }
}, },
}; };
WhileStatement: ast::Stmt = { WhileStatement: ast::Stmt = {
<location:@L> "while" <test:NamedExpressionTest> ":" <body:Suite> <s2:("else" ":" Suite)?> <end_location:@R> => { <location:@L> "while" <test:NamedExpressionTest> ":" <body:Suite> <s2:("else" ":" Suite)?> => {
let orelse = s2.map(|s| s.2).unwrap_or_default(); let orelse = s2.map(|s| s.2).unwrap_or_default();
let end_location = orelse
.last()
.or_else(|| body.last())
.unwrap()
.end_location;
ast::Stmt { ast::Stmt {
custom: (), custom: (),
location, location,
end_location: Some(end_location), end_location,
node: ast::StmtKind::While { node: ast::StmtKind::While {
test: Box::new(test), test: Box::new(test),
body, body,
@ -401,8 +411,14 @@ WhileStatement: ast::Stmt = {
}; };
ForStatement: ast::Stmt = { ForStatement: ast::Stmt = {
<location:@L> <is_async:"async"?> "for" <target:ExpressionList> "in" <iter:TestList> ":" <body:Suite> <s2:("else" ":" Suite)?> <end_location:@R> => { <location:@L> <is_async:"async"?> "for" <target:ExpressionList> "in" <iter:TestList> ":" <body:Suite> <s2:("else" ":" Suite)?> => {
let orelse = s2.map(|s| s.2).unwrap_or_default(); let orelse = s2.map(|s| s.2).unwrap_or_default();
let end_location = orelse
.last()
.or_else(|| body.last())
.unwrap()
.end_location
.unwrap();
let target = Box::new(set_context(target, ast::ExprContext::Store)); let target = Box::new(set_context(target, ast::ExprContext::Store));
let iter = Box::new(iter); let iter = Box::new(iter);
let type_comment = None; let type_comment = None;
@ -419,10 +435,16 @@ TryStatement: ast::Stmt = {
<location:@L> "try" ":" <body:Suite> <handlers:ExceptClause+> <else_suite:("else" ":" Suite)?> <finally:("finally" ":" Suite)?> <end_location:@R> => { <location:@L> "try" ":" <body:Suite> <handlers:ExceptClause+> <else_suite:("else" ":" Suite)?> <finally:("finally" ":" Suite)?> <end_location:@R> => {
let orelse = else_suite.map(|s| s.2).unwrap_or_default(); let orelse = else_suite.map(|s| s.2).unwrap_or_default();
let finalbody = finally.map(|s| s.2).unwrap_or_default(); let finalbody = finally.map(|s| s.2).unwrap_or_default();
let end_location = finalbody
.last()
.map(|last| last.end_location)
.or_else(|| orelse.last().map(|last| last.end_location))
.or_else(|| handlers.last().map(|last| last.end_location))
.unwrap();
ast::Stmt { ast::Stmt {
custom: (), custom: (),
location, location,
end_location: Some(end_location), end_location,
node: ast::StmtKind::Try { node: ast::StmtKind::Try {
body, body,
handlers, handlers,
@ -431,14 +453,15 @@ TryStatement: ast::Stmt = {
}, },
} }
}, },
<location:@L> "try" ":" <body:Suite> <finally:("finally" ":" Suite)> <end_location:@R> => { <location:@L> "try" ":" <body:Suite> <finally:("finally" ":" Suite)> => {
let handlers = vec![]; let handlers = vec![];
let orelse = vec![]; let orelse = vec![];
let finalbody = finally.2; let finalbody = finally.2;
let end_location = finalbody.last().unwrap().end_location;
ast::Stmt { ast::Stmt {
custom: (), custom: (),
location, location,
end_location: Some(end_location), end_location,
node: ast::StmtKind::Try { node: ast::StmtKind::Try {
body, body,
handlers, handlers,
@ -450,7 +473,8 @@ TryStatement: ast::Stmt = {
}; };
ExceptClause: ast::Excepthandler = { ExceptClause: ast::Excepthandler = {
<location:@L> "except" <typ:Test?> ":" <body:Suite> <end_location:@R> => { <location:@L> "except" <typ:Test?> ":" <body:Suite> => {
let end_location = body.last().unwrap().end_location.unwrap();
ast::Excepthandler::new( ast::Excepthandler::new(
location, location,
end_location, end_location,
@ -461,7 +485,8 @@ ExceptClause: ast::Excepthandler = {
}, },
) )
}, },
<location:@L> "except" <x:(Test "as" Identifier)> ":" <body:Suite> <end_location:@R> => { <location:@L> "except" <x:(Test "as" Identifier)> ":" <body:Suite> => {
let end_location = body.last().unwrap().end_location.unwrap();
ast::Excepthandler::new( ast::Excepthandler::new(
location, location,
end_location, end_location,
@ -475,7 +500,8 @@ ExceptClause: ast::Excepthandler = {
}; };
WithStatement: ast::Stmt = { WithStatement: ast::Stmt = {
<location:@L> <is_async:"async"?> "with" <items:WithItems> ":" <body:Suite> <end_location:@R> => { <location:@L> <is_async:"async"?> "with" <items:WithItems> ":" <body:Suite> => {
let end_location = body.last().unwrap().end_location.unwrap();
let type_comment = None; let type_comment = None;
let node = if is_async.is_some() { let node = if is_async.is_some() {
ast::StmtKind::AsyncWith { items, body, type_comment } ast::StmtKind::AsyncWith { items, body, type_comment }
@ -514,9 +540,10 @@ WithItem: ast::Withitem = {
}; };
FuncDef: ast::Stmt = { FuncDef: ast::Stmt = {
<decorator_list:Decorator*> <location:@L> <is_async:"async"?> "def" <name:Identifier> <args:Parameters> <r:("->" Test)?> ":" <body:Suite> <end_location:@R> => { <decorator_list:Decorator*> <location:@L> <is_async:"async"?> "def" <name:Identifier> <args:Parameters> <r:("->" Test)?> ":" <body:Suite> => {
let args = Box::new(args); let args = Box::new(args);
let returns = r.map(|x| Box::new(x.1)); let returns = r.map(|x| Box::new(x.1));
let end_location = body.last().unwrap().end_location.unwrap();
let type_comment = None; let type_comment = None;
let node = if is_async.is_some() { let node = if is_async.is_some() {
ast::StmtKind::AsyncFunctionDef { name, args, body, decorator_list, returns, type_comment } ast::StmtKind::AsyncFunctionDef { name, args, body, decorator_list, returns, type_comment }
@ -668,15 +695,16 @@ KwargParameter<ArgType>: Option<Box<ast::Arg>> = {
}; };
ClassDef: ast::Stmt = { ClassDef: ast::Stmt = {
<decorator_list:Decorator*> <location:@L> "class" <name:Identifier> <a:("(" ArgumentList ")")?> ":" <body:Suite> <end_location:@R> => { <decorator_list:Decorator*> <location:@L> "class" <name:Identifier> <a:("(" ArgumentList ")")?> ":" <body:Suite> => {
let (bases, keywords) = match a { let (bases, keywords) = match a {
Some((_, arg, _)) => (arg.args, arg.keywords), Some((_, arg, _)) => (arg.args, arg.keywords),
None => (vec![], vec![]), None => (vec![], vec![]),
}; };
let end_location = body.last().unwrap().end_location;
ast::Stmt { ast::Stmt {
custom: (), custom: (),
location, location,
end_location: Some(end_location), end_location,
node: ast::StmtKind::ClassDef { node: ast::StmtKind::ClassDef {
name, name,
bases, bases,

View file

@ -173,7 +173,8 @@ class Foo(A, B):
def __init__(self): def __init__(self):
pass pass
def method_with_default(self, arg='default'): def method_with_default(self, arg='default'):
pass"; pass
";
insta::assert_debug_snapshot!(parse_program(source, "<test>").unwrap()); insta::assert_debug_snapshot!(parse_program(source, "<test>").unwrap());
} }

View file

@ -62,8 +62,8 @@ expression: "parse_program(source, \"<test>\").unwrap()"
}, },
end_location: Some( end_location: Some(
Location { Location {
row: 4, row: 3,
column: 1, column: 6,
}, },
), ),
custom: (), custom: (),

View file

@ -79,8 +79,8 @@ expression: parse_ast
}, },
end_location: Some( end_location: Some(
Location { Location {
row: 3, row: 2,
column: 0, column: 10,
}, },
), ),
custom: (), custom: (),

View file

@ -10,8 +10,8 @@ expression: "parse_program(source, \"<test>\").unwrap()"
}, },
end_location: Some( end_location: Some(
Location { Location {
row: 2, row: 1,
column: 0, column: 12,
}, },
), ),
custom: (), custom: (),
@ -66,8 +66,8 @@ expression: "parse_program(source, \"<test>\").unwrap()"
}, },
end_location: Some( end_location: Some(
Location { Location {
row: 3, row: 2,
column: 0, column: 17,
}, },
), ),
custom: (), custom: (),
@ -140,8 +140,8 @@ expression: "parse_program(source, \"<test>\").unwrap()"
}, },
end_location: Some( end_location: Some(
Location { Location {
row: 4, row: 3,
column: 0, column: 15,
}, },
), ),
custom: (), custom: (),
@ -218,8 +218,8 @@ expression: "parse_program(source, \"<test>\").unwrap()"
}, },
end_location: Some( end_location: Some(
Location { Location {
row: 5, row: 4,
column: 0, column: 25,
}, },
), ),
custom: (), custom: (),
@ -332,8 +332,8 @@ expression: "parse_program(source, \"<test>\").unwrap()"
}, },
end_location: Some( end_location: Some(
Location { Location {
row: 6, row: 5,
column: 0, column: 24,
}, },
), ),
custom: (), custom: (),
@ -441,8 +441,8 @@ expression: "parse_program(source, \"<test>\").unwrap()"
}, },
end_location: Some( end_location: Some(
Location { Location {
row: 7, row: 6,
column: 0, column: 29,
}, },
), ),
custom: (), custom: (),
@ -568,8 +568,8 @@ expression: "parse_program(source, \"<test>\").unwrap()"
}, },
end_location: Some( end_location: Some(
Location { Location {
row: 8, row: 7,
column: 0, column: 13,
}, },
), ),
custom: (), custom: (),
@ -622,8 +622,8 @@ expression: "parse_program(source, \"<test>\").unwrap()"
}, },
end_location: Some( end_location: Some(
Location { Location {
row: 9, row: 8,
column: 0, column: 18,
}, },
), ),
custom: (), custom: (),
@ -694,8 +694,8 @@ expression: "parse_program(source, \"<test>\").unwrap()"
}, },
end_location: Some( end_location: Some(
Location { Location {
row: 10, row: 9,
column: 0, column: 14,
}, },
), ),
custom: (), custom: (),
@ -750,8 +750,8 @@ expression: "parse_program(source, \"<test>\").unwrap()"
}, },
end_location: Some( end_location: Some(
Location { Location {
row: 11, row: 10,
column: 0, column: 19,
}, },
), ),
custom: (), custom: (),
@ -824,8 +824,8 @@ expression: "parse_program(source, \"<test>\").unwrap()"
}, },
end_location: Some( end_location: Some(
Location { Location {
row: 12, row: 11,
column: 0, column: 15,
}, },
), ),
custom: (), custom: (),
@ -880,8 +880,8 @@ expression: "parse_program(source, \"<test>\").unwrap()"
}, },
end_location: Some( end_location: Some(
Location { Location {
row: 13, row: 12,
column: 0, column: 20,
}, },
), ),
custom: (), custom: (),
@ -972,8 +972,8 @@ expression: "parse_program(source, \"<test>\").unwrap()"
}, },
end_location: Some( end_location: Some(
Location { Location {
row: 14, row: 13,
column: 0, column: 17,
}, },
), ),
custom: (), custom: (),
@ -1050,8 +1050,8 @@ expression: "parse_program(source, \"<test>\").unwrap()"
}, },
end_location: Some( end_location: Some(
Location { Location {
row: 15, row: 14,
column: 0, column: 22,
}, },
), ),
custom: (), custom: (),
@ -1161,8 +1161,8 @@ expression: "parse_program(source, \"<test>\").unwrap()"
}, },
end_location: Some( end_location: Some(
Location { Location {
row: 16, row: 15,
column: 0, column: 16,
}, },
), ),
custom: (), custom: (),
@ -1249,8 +1249,8 @@ expression: "parse_program(source, \"<test>\").unwrap()"
}, },
end_location: Some( end_location: Some(
Location { Location {
row: 17, row: 16,
column: 0, column: 21,
}, },
), ),
custom: (), custom: (),
@ -1355,8 +1355,8 @@ expression: "parse_program(source, \"<test>\").unwrap()"
}, },
end_location: Some( end_location: Some(
Location { Location {
row: 18, row: 17,
column: 0, column: 18,
}, },
), ),
custom: (), custom: (),
@ -1462,8 +1462,8 @@ expression: "parse_program(source, \"<test>\").unwrap()"
}, },
end_location: Some( end_location: Some(
Location { Location {
row: 19, row: 18,
column: 0, column: 23,
}, },
), ),
custom: (), custom: (),
@ -1587,8 +1587,8 @@ expression: "parse_program(source, \"<test>\").unwrap()"
}, },
end_location: Some( end_location: Some(
Location { Location {
row: 20, row: 19,
column: 0, column: 19,
}, },
), ),
custom: (), custom: (),
@ -1675,8 +1675,8 @@ expression: "parse_program(source, \"<test>\").unwrap()"
}, },
end_location: Some( end_location: Some(
Location { Location {
row: 21, row: 20,
column: 0, column: 24,
}, },
), ),
custom: (), custom: (),
@ -1781,8 +1781,8 @@ expression: "parse_program(source, \"<test>\").unwrap()"
}, },
end_location: Some( end_location: Some(
Location { Location {
row: 22, row: 21,
column: 0, column: 27,
}, },
), ),
custom: (), custom: (),
@ -1923,8 +1923,8 @@ expression: "parse_program(source, \"<test>\").unwrap()"
}, },
end_location: Some( end_location: Some(
Location { Location {
row: 23, row: 22,
column: 0, column: 32,
}, },
), ),
custom: (), custom: (),