Improve PLW0120 range (#1500)

src/ast/helpers.rs

@@ -426,6 +426,32 @@ pub fn excepthandler_name_range(
     }
 }
 
+/// Return the `Range` of `else` in `For`, `AsyncFor`, and `While` statements.
+pub fn else_range(stmt: &Stmt, locator: &SourceCodeLocator) -> Option<Range> {
+    match &stmt.node {
+        StmtKind::For { body, orelse, .. }
+        | StmtKind::AsyncFor { body, orelse, .. }
+        | StmtKind::While { body, orelse, .. }
+            if !orelse.is_empty() =>
+        {
+            let body_end = body.last().unwrap().end_location.unwrap();
+            let contents = locator.slice_source_code_range(&Range {
+                location: body_end,
+                end_location: orelse[0].location,
+            });
+            let range = lexer::make_tokenizer_located(&contents, body_end)
+                .flatten()
+                .find(|(_, kind, _)| matches!(kind, Tok::Else))
+                .map(|(location, _, end_location)| Range {
+                    location,
+                    end_location,
+                });
+            range
+        }
+        _ => None,
+    }
+}
+
 /// 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, locator: &SourceCodeLocator) -> bool {
@@ -466,7 +492,9 @@ mod tests {
     use rustpython_ast::Location;
     use rustpython_parser::parser;
 
-    use crate::ast::helpers::{identifier_range, match_module_member, match_trailing_content};
+    use crate::ast::helpers::{
+        else_range, identifier_range, match_module_member, match_trailing_content,
+    };
     use crate::ast::types::Range;
     use crate::source_code_locator::SourceCodeLocator;
 
@@ -738,4 +766,24 @@ class Class():
 
         Ok(())
     }
+
+    #[test]
+    fn test_else_range() -> Result<()> {
+        let contents = r#"
+for x in y:
+    pass
+else:
+    pass
+"#
+        .trim();
+        let program = parser::parse_program(contents, "<filename>")?;
+        let stmt = program.first().unwrap();
+        let locator = SourceCodeLocator::new(contents);
+        let range = 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(())
+    }
 }

src/pylint/plugins/useless_else_on_loop.rs

@@ -1,6 +1,6 @@
 use rustpython_ast::{ExcepthandlerKind, Stmt, StmtKind};
 
-use crate::ast::types::Range;
+use crate::ast::helpers;
 use crate::checkers::ast::Checker;
 use crate::checks::CheckKind;
 use crate::Check;
@@ -35,7 +35,7 @@ pub fn useless_else_on_loop(checker: &mut Checker, stmt: &Stmt, body: &[Stmt], o
     if !orelse.is_empty() && !loop_exits_early(body) {
         checker.add_check(Check::new(
             CheckKind::UselessElseOnLoop,
-            Range::from_located(stmt),
+            helpers::else_range(stmt, checker.locator).unwrap(),
         ));
     }
 }

src/pylint/snapshots/ruff__pylint__tests__PLW0120_useless_else_on_loop.py.snap

@@ -4,65 +4,65 @@ expression: checks
 ---
 - kind: UselessElseOnLoop
   location:
-    row: 6
+    row: 9
     column: 4
   end_location:
-    row: 10
-    column: 31
+    row: 9
+    column: 8
   fix: ~
   parent: ~
 - kind: UselessElseOnLoop
   location:
-    row: 16
+    row: 18
     column: 4
   end_location:
-    row: 19
-    column: 31
+    row: 18
+    column: 8
   fix: ~
   parent: ~
 - kind: UselessElseOnLoop
   location:
-    row: 23
+    row: 30
     column: 0
   end_location:
-    row: 31
-    column: 21
+    row: 30
+    column: 4
   fix: ~
   parent: ~
 - kind: UselessElseOnLoop
   location:
-    row: 34
+    row: 37
     column: 0
   end_location:
-    row: 38
-    column: 21
+    row: 37
+    column: 4
   fix: ~
   parent: ~
 - kind: UselessElseOnLoop
   location:
-    row: 40
+    row: 42
     column: 0
   end_location:
-    row: 45
-    column: 13
+    row: 42
+    column: 4
   fix: ~
   parent: ~
 - kind: UselessElseOnLoop
   location:
-    row: 81
+    row: 88
     column: 4
   end_location:
-    row: 89
-    column: 19
+    row: 88
+    column: 8
   fix: ~
   parent: ~
 - kind: UselessElseOnLoop
   location:
-    row: 96
+    row: 98
     column: 8
   end_location:
-    row: 100
-    column: 21
+    row: 98
+    column: 12
   fix: ~
   parent: ~