diff --git a/src/ast/helpers.rs b/src/ast/helpers.rs index 21be4bce03..e6eaa32e84 100644 --- a/src/ast/helpers.rs +++ b/src/ast/helpers.rs @@ -614,6 +614,36 @@ pub fn first_colon_range(range: Range, locator: &Locator) -> Option { range } +/// Return the `Range` of the first `Elif` or `Else` token in an `If` statement. +pub fn elif_else_range(stmt: &Stmt, locator: &Locator) -> Option { + let StmtKind::If { body, orelse, .. } = &stmt.node else { + return None; + }; + + let start = body + .last() + .expect("Expected body to be non-empty") + .end_location + .unwrap(); + let end = match &orelse[..] { + [Stmt { + node: StmtKind::If { test, .. }, + .. + }] => test.location, + [stmt, ..] => stmt.location, + _ => return None, + }; + let contents = locator.slice_source_code_range(&Range::new(start, end)); + let range = lexer::make_tokenizer_located(&contents, start) + .flatten() + .find(|(_, kind, _)| matches!(kind, Tok::Elif | Tok::Else)) + .map(|(location, _, end_location)| Range { + location, + end_location, + }); + range +} + /// Return `true` if a `Stmt` appears to be part of a multi-statement line, with /// other statements preceding it. pub fn preceded_by_continuation(stmt: &Stmt, indexer: &Indexer) -> bool { @@ -709,7 +739,7 @@ mod tests { use rustpython_parser::parser; use crate::ast::helpers::{ - else_range, first_colon_range, identifier_range, match_trailing_content, + elif_else_range, else_range, first_colon_range, identifier_range, match_trailing_content, }; use crate::ast::types::Range; use crate::source_code::Locator; @@ -856,4 +886,39 @@ else: assert_eq!(range.end_location.row(), 1); assert_eq!(range.end_location.column(), 7); } + + #[test] + fn test_elif_else_range() -> Result<()> { + let contents = " +if a: + ... +elif b: + ... +" + .trim_start(); + let program = parser::parse_program(contents, "")?; + let stmt = program.first().unwrap(); + let locator = Locator::new(contents); + let range = elif_else_range(stmt, &locator).unwrap(); + assert_eq!(range.location.row(), 3); + assert_eq!(range.location.column(), 0); + assert_eq!(range.end_location.row(), 3); + assert_eq!(range.end_location.column(), 4); + let contents = " +if a: + ... +else: + ... +" + .trim_start(); + let program = parser::parse_program(contents, "")?; + let stmt = program.first().unwrap(); + let locator = Locator::new(contents); + let range = elif_else_range(stmt, &locator).unwrap(); + assert_eq!(range.location.row(), 3); + assert_eq!(range.location.column(), 0); + assert_eq!(range.end_location.row(), 3); + assert_eq!(range.end_location.column(), 4); + Ok(()) + } } diff --git a/src/rules/flake8_return/rules.rs b/src/rules/flake8_return/rules.rs index 58da12dfc7..214888e016 100644 --- a/src/rules/flake8_return/rules.rs +++ b/src/rules/flake8_return/rules.rs @@ -3,6 +3,7 @@ use rustpython_ast::{Constant, Expr, ExprKind, Location, Stmt, StmtKind}; use super::helpers::result_exists; use super::visitor::{ReturnVisitor, Stack}; +use crate::ast::helpers::elif_else_range; use crate::ast::types::Range; use crate::ast::visitor::Visitor; use crate::ast::whitespace::indentation; @@ -228,7 +229,8 @@ fn superfluous_else_node(checker: &mut Checker, stmt: &Stmt, branch: Branch) -> if checker.settings.enabled.contains(&RuleCode::RET505) { checker.diagnostics.push(Diagnostic::new( violations::SuperfluousElseReturn(branch), - Range::from_located(stmt), + elif_else_range(stmt, checker.locator) + .unwrap_or_else(|| Range::from_located(stmt)), )); } return true; @@ -237,7 +239,8 @@ fn superfluous_else_node(checker: &mut Checker, stmt: &Stmt, branch: Branch) -> if checker.settings.enabled.contains(&RuleCode::RET508) { checker.diagnostics.push(Diagnostic::new( violations::SuperfluousElseBreak(branch), - Range::from_located(stmt), + elif_else_range(stmt, checker.locator) + .unwrap_or_else(|| Range::from_located(stmt)), )); } return true; @@ -246,7 +249,8 @@ fn superfluous_else_node(checker: &mut Checker, stmt: &Stmt, branch: Branch) -> if checker.settings.enabled.contains(&RuleCode::RET506) { checker.diagnostics.push(Diagnostic::new( violations::SuperfluousElseRaise(branch), - Range::from_located(stmt), + elif_else_range(stmt, checker.locator) + .unwrap_or_else(|| Range::from_located(stmt)), )); } return true; @@ -255,7 +259,8 @@ fn superfluous_else_node(checker: &mut Checker, stmt: &Stmt, branch: Branch) -> if checker.settings.enabled.contains(&RuleCode::RET507) { checker.diagnostics.push(Diagnostic::new( violations::SuperfluousElseContinue(branch), - Range::from_located(stmt), + elif_else_range(stmt, checker.locator) + .unwrap_or_else(|| Range::from_located(stmt)), )); } return true; diff --git a/src/rules/flake8_return/snapshots/ruff__rules__flake8_return__tests__RET505_RET505.py.snap b/src/rules/flake8_return/snapshots/ruff__rules__flake8_return__tests__RET505_RET505.py.snap index 355a3af3aa..33fe112d68 100644 --- a/src/rules/flake8_return/snapshots/ruff__rules__flake8_return__tests__RET505_RET505.py.snap +++ b/src/rules/flake8_return/snapshots/ruff__rules__flake8_return__tests__RET505_RET505.py.snap @@ -5,81 +5,81 @@ expression: diagnostics - kind: SuperfluousElseReturn: Elif location: - row: 5 + row: 8 column: 4 end_location: - row: 13 - column: 16 + row: 8 + column: 8 fix: ~ parent: ~ - kind: SuperfluousElseReturn: Elif location: - row: 17 + row: 23 column: 4 end_location: - row: 26 - column: 13 + row: 23 + column: 8 fix: ~ parent: ~ - kind: SuperfluousElseReturn: Elif location: - row: 38 + row: 41 column: 4 end_location: - row: 46 - column: 16 + row: 41 + column: 8 fix: ~ parent: ~ - kind: SuperfluousElseReturn: Else location: - row: 50 + row: 53 column: 4 end_location: - row: 55 - column: 16 + row: 53 + column: 8 fix: ~ parent: ~ - kind: SuperfluousElseReturn: Else location: - row: 61 + row: 64 column: 8 end_location: - row: 66 - column: 20 + row: 64 + column: 12 fix: ~ parent: ~ - kind: SuperfluousElseReturn: Else location: - row: 73 + row: 79 column: 4 end_location: - row: 80 - column: 13 + row: 79 + column: 8 fix: ~ parent: ~ - kind: SuperfluousElseReturn: Else location: - row: 86 + row: 89 column: 8 end_location: - row: 90 - column: 17 + row: 89 + column: 12 fix: ~ parent: ~ - kind: SuperfluousElseReturn: Else location: - row: 97 + row: 99 column: 4 end_location: - row: 103 - column: 23 + row: 99 + column: 8 fix: ~ parent: ~ diff --git a/src/rules/flake8_return/snapshots/ruff__rules__flake8_return__tests__RET506_RET506.py.snap b/src/rules/flake8_return/snapshots/ruff__rules__flake8_return__tests__RET506_RET506.py.snap index a67eb39898..e18db00168 100644 --- a/src/rules/flake8_return/snapshots/ruff__rules__flake8_return__tests__RET506_RET506.py.snap +++ b/src/rules/flake8_return/snapshots/ruff__rules__flake8_return__tests__RET506_RET506.py.snap @@ -5,71 +5,71 @@ expression: diagnostics - kind: SuperfluousElseRaise: Elif location: - row: 5 + row: 8 column: 4 end_location: - row: 13 - column: 26 + row: 8 + column: 8 fix: ~ parent: ~ - kind: SuperfluousElseRaise: Elif location: - row: 17 + row: 23 column: 4 end_location: - row: 26 - column: 13 + row: 23 + column: 8 fix: ~ parent: ~ - kind: SuperfluousElseRaise: Else location: - row: 31 + row: 34 column: 4 end_location: - row: 36 - column: 26 + row: 34 + column: 8 fix: ~ parent: ~ - kind: SuperfluousElseRaise: Else location: - row: 42 + row: 45 column: 8 end_location: - row: 47 - column: 30 + row: 45 + column: 12 fix: ~ parent: ~ - kind: SuperfluousElseRaise: Else location: - row: 54 + row: 60 column: 4 end_location: - row: 61 - column: 13 + row: 60 + column: 8 fix: ~ parent: ~ - kind: SuperfluousElseRaise: Else location: - row: 67 + row: 70 column: 8 end_location: - row: 71 - column: 17 + row: 70 + column: 12 fix: ~ parent: ~ - kind: SuperfluousElseRaise: Else location: - row: 78 + row: 80 column: 4 end_location: - row: 84 - column: 33 + row: 80 + column: 8 fix: ~ parent: ~ diff --git a/src/rules/flake8_return/snapshots/ruff__rules__flake8_return__tests__RET507_RET507.py.snap b/src/rules/flake8_return/snapshots/ruff__rules__flake8_return__tests__RET507_RET507.py.snap index 6583d720e3..c9df517a21 100644 --- a/src/rules/flake8_return/snapshots/ruff__rules__flake8_return__tests__RET507_RET507.py.snap +++ b/src/rules/flake8_return/snapshots/ruff__rules__flake8_return__tests__RET507_RET507.py.snap @@ -5,71 +5,71 @@ expression: diagnostics - kind: SuperfluousElseContinue: Elif location: - row: 6 + row: 8 column: 8 end_location: - row: 11 - column: 17 + row: 8 + column: 12 fix: ~ parent: ~ - kind: SuperfluousElseContinue: Elif location: - row: 16 + row: 22 column: 8 end_location: - row: 25 - column: 17 + row: 22 + column: 12 fix: ~ parent: ~ - kind: SuperfluousElseContinue: Else location: - row: 34 + row: 36 column: 8 end_location: - row: 37 - column: 17 + row: 36 + column: 12 fix: ~ parent: ~ - kind: SuperfluousElseContinue: Else location: - row: 44 + row: 47 column: 12 end_location: - row: 49 - column: 24 + row: 47 + column: 16 fix: ~ parent: ~ - kind: SuperfluousElseContinue: Else location: - row: 57 + row: 63 column: 8 end_location: - row: 64 - column: 17 + row: 63 + column: 12 fix: ~ parent: ~ - kind: SuperfluousElseContinue: Else location: - row: 71 + row: 74 column: 12 end_location: - row: 75 - column: 21 + row: 74 + column: 16 fix: ~ parent: ~ - kind: SuperfluousElseContinue: Else location: - row: 83 + row: 85 column: 8 end_location: - row: 89 - column: 24 + row: 85 + column: 12 fix: ~ parent: ~ diff --git a/src/rules/flake8_return/snapshots/ruff__rules__flake8_return__tests__RET508_RET508.py.snap b/src/rules/flake8_return/snapshots/ruff__rules__flake8_return__tests__RET508_RET508.py.snap index 069d3d0706..d9afe74869 100644 --- a/src/rules/flake8_return/snapshots/ruff__rules__flake8_return__tests__RET508_RET508.py.snap +++ b/src/rules/flake8_return/snapshots/ruff__rules__flake8_return__tests__RET508_RET508.py.snap @@ -5,71 +5,71 @@ expression: diagnostics - kind: SuperfluousElseBreak: Elif location: - row: 6 + row: 8 column: 8 end_location: - row: 11 - column: 17 + row: 8 + column: 12 fix: ~ parent: ~ - kind: SuperfluousElseBreak: Elif location: - row: 16 + row: 22 column: 8 end_location: - row: 25 - column: 17 + row: 22 + column: 12 fix: ~ parent: ~ - kind: SuperfluousElseBreak: Else location: - row: 31 + row: 33 column: 8 end_location: - row: 34 - column: 17 + row: 33 + column: 12 fix: ~ parent: ~ - kind: SuperfluousElseBreak: Else location: - row: 41 + row: 44 column: 12 end_location: - row: 46 - column: 21 + row: 44 + column: 16 fix: ~ parent: ~ - kind: SuperfluousElseBreak: Else location: - row: 54 + row: 60 column: 8 end_location: - row: 61 - column: 17 + row: 60 + column: 12 fix: ~ parent: ~ - kind: SuperfluousElseBreak: Else location: - row: 68 + row: 71 column: 12 end_location: - row: 72 - column: 21 + row: 71 + column: 16 fix: ~ parent: ~ - kind: SuperfluousElseBreak: Else location: - row: 80 + row: 82 column: 8 end_location: - row: 86 - column: 21 + row: 82 + column: 12 fix: ~ parent: ~