Fix F841 (UnusedVariable) range in except handler (#1367)

astral-sh · Dec 25, 2022 · 10f75c9 · 10f75c9
1 parent cc21104
commit 10f75c9
Show file tree

Hide file tree

Showing 4 changed files with 32 additions and 11 deletions.
diff --git a/src/ast/helpers.rs b/src/ast/helpers.rs
@@ -1,3 +1,4 @@
+use itertools::Itertools;
 use log::error;
 use once_cell::sync::Lazy;
 use regex::Regex;
@@ -396,6 +397,23 @@ pub fn identifier_range(stmt: &Stmt, locator: &SourceCodeLocator) -> Range {
     Range::from_located(stmt)
 }
 
+/// Return the `Range` of `name` in `Excepthandler`.
+pub fn excepthandler_name_range(
+    handler: &Excepthandler,
+    locator: &SourceCodeLocator,
+) -> Option<Range> {
+    let contents = locator.slice_source_code_range(&Range::from_located(handler));
+    let range = lexer::make_tokenizer(&contents)
+        .flatten()
+        .tuple_windows()
+        .find(|(tok, next_tok)| matches!(tok.1, Tok::As) && matches!(next_tok.1, Tok::Name { .. }))
+        .map(|((..), (start, _, end))| Range {
+            location: to_absolute(start, handler.location),
+            end_location: to_absolute(end, handler.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, locator: &SourceCodeLocator) -> bool {

diff --git a/src/checkers/ast.rs b/src/checkers/ast.rs
@@ -2664,12 +2664,15 @@ where
 
                         self.check_builtin_shadowing(name, excepthandler, false);
 
+                        let name_range =
+                            helpers::excepthandler_name_range(excepthandler, self.locator).unwrap();
+
                         if self.current_scope().values.contains_key(&name.as_str()) {
                             self.handle_node_store(
                                 name,
                                 &Expr::new(
-                                    excepthandler.location,
-                                    excepthandler.end_location.unwrap(),
+                                    name_range.location,
+                                    name_range.end_location,
                                     ExprKind::Name {
                                         id: name.to_string(),
                                         ctx: ExprContext::Store,
@@ -2682,8 +2685,8 @@ where
                         self.handle_node_store(
                             name,
                             &Expr::new(
-                                excepthandler.location,
-                                excepthandler.end_location.unwrap(),
+                                name_range.location,
+                                name_range.end_location,
                                 ExprKind::Name {
                                     id: name.to_string(),
                                     ctx: ExprContext::Store,
@@ -2702,7 +2705,7 @@ where
                                 if self.settings.enabled.contains(&CheckCode::F841) {
                                     self.add_check(Check::new(
                                         CheckKind::UnusedVariable(name.to_string()),
-                                        Range::from_located(excepthandler),
+                                        name_range,
                                     ));
                                 }
                             }

diff --git a/src/pyflakes/snapshots/ruff__pyflakes__tests__F841_F841_0.py.snap b/src/pyflakes/snapshots/ruff__pyflakes__tests__F841_F841_0.py.snap
@@ -6,10 +6,10 @@ expression: checks
     UnusedVariable: e
   location:
     row: 3
-    column: 0
+    column: 21
   end_location:
-    row: 4
-    column: 8
+    row: 3
+    column: 22
   fix: ~
 - kind:
     UnusedVariable: z

diff --git a/src/pyflakes/snapshots/ruff__pyflakes__tests__f841_dummy_variable_rgx.snap b/src/pyflakes/snapshots/ruff__pyflakes__tests__f841_dummy_variable_rgx.snap
@@ -6,10 +6,10 @@ expression: checks
     UnusedVariable: e
   location:
     row: 3
-    column: 0
+    column: 21
   end_location:
-    row: 4
-    column: 8
+    row: 3
+    column: 22
   fix: ~
 - kind:
     UnusedVariable: foo