Store body statements instead of index into class_def.body vec

astral-sh · Feb 22, 2024 · 45d9f7f · 45d9f7f
1 parent 309527d
commit 45d9f7f
Show file tree

Hide file tree

Showing 4 changed files with 11 additions and 16 deletions.
diff --git a/crates/ruff_linter/resources/test/fixtures/pylint/no_method_decorator.py b/crates/ruff_linter/resources/test/fixtures/pylint/no_method_decorator.py
@@ -22,7 +22,7 @@ class Class:
     def class_method(cls):
         pass
 
-    class_method = classmethod(class_method);
+    class_method = classmethod(class_method);another_statement
 
     def static_method():
         pass

diff --git a/crates/ruff_linter/src/rules/pylint/rules/no_method_decorator.rs b/crates/ruff_linter/src/rules/pylint/rules/no_method_decorator.rs
@@ -101,15 +101,15 @@ fn get_undecorated_methods(checker: &mut Checker, class_stmt: &Stmt, method_type
         return;
     };
 
-    let mut explicit_decorator_calls: HashMap<String, usize> = HashMap::default();
+    let mut explicit_decorator_calls: HashMap<String, &Stmt> = HashMap::default();
 
     let (method_name, diagnostic_type): (&str, DiagnosticKind) = match method_type {
         MethodType::Classmethod => ("classmethod", NoClassmethodDecorator.into()),
         MethodType::Staticmethod => ("staticmethod", NoStaticmethodDecorator.into()),
     };
 
     // gather all explicit *method calls
-    for (i, stmt) in class_def.body.iter().enumerate() {
+    for stmt in class_def.body.iter() {
         if let Stmt::Assign(ast::StmtAssign { targets, value, .. }) = stmt {
             if let Expr::Call(ast::ExprCall {
                 func, arguments, ..
@@ -132,7 +132,7 @@ fn get_undecorated_methods(checker: &mut Checker, class_stmt: &Stmt, method_type
 
                         if let Expr::Name(ast::ExprName { id, .. }) = &arguments.args[0] {
                             if target_name == *id {
-                                explicit_decorator_calls.insert(id.clone(), i);
+                                explicit_decorator_calls.insert(id.clone(), stmt);
                             }
                         };
                     }
@@ -152,7 +152,7 @@ fn get_undecorated_methods(checker: &mut Checker, class_stmt: &Stmt, method_type
             ..
         }) = stmt
         {
-            if !explicit_decorator_calls.contains_key(name.as_str()) {
+            let Some(decorator_call_statement) = explicit_decorator_calls.get(name.as_str()) else {
                 continue;
             };
 
@@ -178,14 +178,13 @@ fn get_undecorated_methods(checker: &mut Checker, class_stmt: &Stmt, method_type
 
             match indentation {
                 Some(indentation) => {
-                    let i = explicit_decorator_calls[name.as_str()];
                     diagnostic.set_fix(Fix::safe_edits(
                         Edit::insertion(
                             format!("@{method_name}\n{indentation}"),
                             stmt.range().start(),
                         ),
                         [fix::edits::delete_stmt(
-                            &class_def.body[i],
+                            decorator_call_statement,
                             Some(class_stmt),
                             checker.locator(),
                             checker.indexer(),

diff --git a/...s/pylint/snapshots/ruff_linter__rules__pylint__tests__PLR0202_no_method_decorator.py.snap b/...s/pylint/snapshots/ruff_linter__rules__pylint__tests__PLR0202_no_method_decorator.py.snap
@@ -43,10 +43,8 @@ no_method_decorator.py:22:5: PLR0202 [*] Class method defined without decorator
 22 23 |     def class_method(cls):
 23 24 |         pass
 24 25 | 
-25    |-    class_method = classmethod(class_method);
-   26 |+    
+25    |-    class_method = classmethod(class_method);another_statement
+   26 |+    another_statement
 26 27 | 
 27 28 |     def static_method():
 28 29 |         pass
-
-
diff --git a/...s/pylint/snapshots/ruff_linter__rules__pylint__tests__PLR0203_no_method_decorator.py.snap b/...s/pylint/snapshots/ruff_linter__rules__pylint__tests__PLR0203_no_method_decorator.py.snap
@@ -28,7 +28,7 @@ no_method_decorator.py:15:5: PLR0203 [*] Static method defined without decorator
 
 no_method_decorator.py:27:5: PLR0203 [*] Static method defined without decorator
    |
-25 |     class_method = classmethod(class_method);
+25 |     class_method = classmethod(class_method);another_statement
 26 | 
 27 |     def static_method():
    |      PLR0203
@@ -38,13 +38,11 @@ no_method_decorator.py:27:5: PLR0203 [*] Static method defined without decorator
 
 ℹ Safe fix
 24 24 | 
-25 25 |     class_method = classmethod(class_method);
+25 25 |     class_method = classmethod(class_method);another_statement
 26 26 | 
    27 |+    @staticmethod
 27 28 |     def static_method():
 28 29 |         pass
 29 30 | 
 30    |-    static_method = staticmethod(static_method);
-   31 |+    
-
-
+   31 |+