Avoid removing trailing comments on pass statements

resources/test/fixtures/flake8_pie/PIE790.py

@@ -1,12 +1,12 @@
 class Foo:
     """buzz"""
 
-    pass  # PIE790
+    pass
 
 
 if foo:
     """foo"""
-    pass  # PIE790
+    pass
 
 
 def multi_statement() -> None:
@@ -18,51 +18,51 @@ def multi_statement() -> None:
     pass
 else:
     """bar"""
-    pass  # PIE790
+    pass
 
 
 while True:
     pass
 else:
     """bar"""
-    pass  # PIE790
+    pass
 
 
 for _ in range(10):
     pass
 else:
     """bar"""
-    pass  # PIE790
+    pass
 
 
 async for _ in range(10):
     pass
 else:
     """bar"""
-    pass  # PIE790
+    pass
 
 
 def foo() -> None:
     """
     buzz
     """
 
-    pass  # PIE790
+    pass
 
 
 async def foo():
     """
     buzz
     """
 
-    pass  # PIE790
+    pass
 
 
 try:
     """
     buzz
     """
-    pass  # PIE790
+    pass
 except ValueError:
     pass
 
@@ -71,29 +71,34 @@ async def foo():
     bar()
 except ValueError:
     """bar"""
-    pass  # PIE790
+    pass
 
 
 for _ in range(10):
     """buzz"""
-    pass  # PIE790
+    pass
 
 async for _ in range(10):
     """buzz"""
-    pass  # PIE790
+    pass
 
 while cond:
     """buzz"""
-    pass  # PIE790
+    pass
 
 
 with bar:
     """buzz"""
-    pass  # PIE790
+    pass
 
 async with bar:
     """buzz"""
-    pass  # PIE790
+    pass
+
+
+def foo() -> None:
+    """buzz"""
+    pass  # bar
 
 
 class Foo:

src/ast/helpers.rs

@@ -665,6 +665,24 @@ pub fn match_trailing_content(stmt: &Stmt, locator: &Locator) -> bool {
     false
 }
 
+/// If a `Stmt` has a trailing comment, return the index of the hash.
+pub fn match_trailing_comment(stmt: &Stmt, locator: &Locator) -> Option<usize> {
+    let range = Range::new(
+        stmt.end_location.unwrap(),
+        Location::new(stmt.end_location.unwrap().row() + 1, 0),
+    );
+    let suffix = locator.slice_source_code_range(&range);
+    for (i, char) in suffix.chars().enumerate() {
+        if char == '#' {
+            return Some(i);
+        }
+        if !char.is_whitespace() {
+            return None;
+        }
+    }
+    None
+}
+
 /// Return the number of trailing empty lines following a statement.
 pub fn count_trailing_lines(stmt: &Stmt, locator: &Locator) -> usize {
     let suffix =

src/rules/flake8_pie/rules.rs

@@ -3,11 +3,12 @@ use rustc_hash::FxHashSet;
 use rustpython_ast::{Constant, Expr, ExprKind, Keyword, Stmt, StmtKind};
 
 use crate::ast::comparable::ComparableExpr;
-use crate::ast::helpers::unparse_expr;
+use crate::ast::helpers::{match_trailing_comment, unparse_expr};
 use crate::ast::types::{Range, RefEquality};
 use crate::autofix::helpers::delete_stmt;
 use crate::checkers::ast::Checker;
 use crate::fix::Fix;
+use crate::message::Location;
 use crate::python::identifiers::is_identifier;
 use crate::registry::{Diagnostic, Rule};
 use crate::violations;
@@ -35,19 +36,29 @@ pub fn no_unnecessary_pass(checker: &mut Checker, body: &[Stmt]) {
                     Range::from_located(pass_stmt),
                 );
                 if checker.patch(&Rule::NoUnnecessaryPass) {
-                    match delete_stmt(
-                        pass_stmt,
-                        None,
-                        &[],
-                        checker.locator,
-                        checker.indexer,
-                        checker.stylist,
-                    ) {
-                        Ok(fix) => {
-                            diagnostic.amend(fix);
-                        }
-                        Err(e) => {
-                            error!("Failed to delete `pass` statement: {}", e);
+                    if let Some(index) = match_trailing_comment(pass_stmt, checker.locator) {
+                        diagnostic.amend(Fix::deletion(
+                            pass_stmt.location,
+                            Location::new(
+                                pass_stmt.end_location.unwrap().row(),
+                                pass_stmt.end_location.unwrap().column() + index,
+                            ),
+                        ));
+                    } else {
+                        match delete_stmt(
+                            pass_stmt,
+                            None,
+                            &[],
+                            checker.locator,
+                            checker.indexer,
+                            checker.stylist,
+                        ) {
+                            Ok(fix) => {
+                                diagnostic.amend(fix);
+                            }
+                            Err(e) => {
+                                error!("Failed to delete `pass` statement: {}", e);
+                            }
                         }
                     }
                 }

src/rules/flake8_pie/snapshots/ruff__rules__flake8_pie__tests__PIE790_PIE790.py.snap

@@ -290,4 +290,22 @@ expression: diagnostics
       row: 97
       column: 0
   parent: ~
+- kind:
+    NoUnnecessaryPass: ~
+  location:
+    row: 101
+    column: 4
+  end_location:
+    row: 101
+    column: 8
+  fix:
+    content:
+      - ""
+    location:
+      row: 101
+      column: 4
+    end_location:
+      row: 101
+      column: 10
+  parent: ~