[refurb] Preserve comments (FURB171)

crates/ruff_linter/resources/test/fixtures/refurb/FURB171.py

@@ -46,3 +46,62 @@
 
 if 1 in {*[1]}:
     pass
+
+
+# https://github.com/astral-sh/ruff/issues/10063
+_ = a in (
+    # Foo
+    b,
+)
+
+_ = a in (  # Foo1
+    (  # Foo2
+    # Foo3
+        (  # Tuple
+            (  # Bar
+       (b
+        # Bar
+        )
+  )
+       # Foo4
+            # Foo5
+     ,
+       )
+        # Foo6
+    )
+)
+
+foo = (
+    lorem()
+        .ipsum()
+        .dolor(lambda sit: sit in (
+            # Foo1
+            # Foo2
+            amet,
+        ))
+)
+
+foo = lorem() \
+    .ipsum() \
+    .dolor(lambda sit: sit in (
+        # Foo1
+        # Foo2
+        amet,
+    ))
+
+def _():
+    if foo not \
+        in [
+        # Before
+        bar
+        # After
+    ]: ...
+
+def _():
+    if foo not \
+        in [
+        # Before
+        bar
+        # After
+    ] and \
+        0 < 1: ...

crates/ruff_linter/src/rules/refurb/rules/single_item_membership_test.rs

@@ -1,11 +1,13 @@
-use ruff_diagnostics::{Applicability, Diagnostic, Edit, Fix, FixAvailability, Violation};
+use ruff_diagnostics::{AlwaysFixableViolation, Applicability, Diagnostic, Edit, Fix};
 use ruff_macros::{derive_message_formats, ViolationMetadata};
-use ruff_python_ast::helpers::generate_comparison;
+use ruff_python_ast::parenthesize::parenthesized_range;
 use ruff_python_ast::{self as ast, CmpOp, Expr, ExprStringLiteral};
-use ruff_text_size::Ranged;
+use ruff_python_parser::TokenKind;
+use ruff_source_file::LineRanges;
+use ruff_text_size::{Ranged, TextRange, TextSize};
+use smallvec::{smallvec, SmallVec};
 
 use crate::checkers::ast::Checker;
-use crate::fix::edits::pad;
 
 /// ## What it does
 /// Checks for membership tests against single-item containers.
@@ -39,19 +41,17 @@ pub(crate) struct SingleItemMembershipTest {
     membership_test: MembershipTest,
 }
 
-impl Violation for SingleItemMembershipTest {
-    const FIX_AVAILABILITY: FixAvailability = FixAvailability::Sometimes;
-
+impl AlwaysFixableViolation for SingleItemMembershipTest {
     #[derive_message_formats]
     fn message(&self) -> String {
         "Membership test against single-item container".to_string()
     }
 
-    fn fix_title(&self) -> Option<String> {
+    fn fix_title(&self) -> String {
         let SingleItemMembershipTest { membership_test } = self;
         match membership_test {
-            MembershipTest::In => Some("Convert to equality test".to_string()),
-            MembershipTest::NotIn => Some("Convert to inequality test".to_string()),
+            MembershipTest::In { .. } => "Convert to equality test".to_string(),
+            MembershipTest::NotIn { .. } => "Convert to inequality test".to_string(),
         }
     }
 }
@@ -68,10 +68,32 @@ pub(crate) fn single_item_membership_test(
         return;
     };
 
+    let tokens = checker.tokens();
+    let find_token_after = |offset: TextSize, kind: TokenKind| {
+        tokens
+            .after(offset)
+            .iter()
+            .find(|token| token.kind() == kind)
+            .unwrap()
+    };
+
     // Ensure that the comparison is a membership test.
     let membership_test = match op {
-        CmpOp::In => MembershipTest::In,
-        CmpOp::NotIn => MembershipTest::NotIn,
+        CmpOp::In => {
+            let in_token = find_token_after(left.end(), TokenKind::In);
+
+            MembershipTest::In {
+                range: in_token.range(),
+            }
+        }
+        CmpOp::NotIn => {
+            let not_token = find_token_after(left.end(), TokenKind::Not);
+            let in_token = find_token_after(not_token.end(), TokenKind::In);
+
+            MembershipTest::NotIn {
+                range: TextRange::new(not_token.start(), in_token.end()),
+            }
+        }
         _ => return,
     };
 
@@ -81,30 +103,7 @@ pub(crate) fn single_item_membership_test(
     };
 
     let diagnostic = Diagnostic::new(SingleItemMembershipTest { membership_test }, expr.range());
-
-    let edit = Edit::range_replacement(
-        pad(
-            generate_comparison(
-                left,
-                &[membership_test.replacement_op()],
-                &[item.clone()],
-                expr.into(),
-                checker.comment_ranges(),
-                checker.source(),
-            ),
-            expr.range(),
-            checker.locator(),
-        ),
-        expr.range(),
-    );
-
-    let applicability = if right.is_string_literal_expr() {
-        Applicability::Unsafe
-    } else {
-        Applicability::Safe
-    };
-
-    let fix = Fix::applicable_edit(edit, applicability);
+    let fix = replace_with_comparison(membership_test, right, item, checker);
 
     checker.diagnostics.push(diagnostic.with_fix(fix));
 }
@@ -129,20 +128,103 @@ fn single_item(expr: &Expr) -> Option<&Expr> {
     }
 }
 
+fn replace_with_comparison(
+    membership_test: MembershipTest,
+    iterable: &Expr,
+    item: &Expr,
+    checker: &Checker,
+) -> Fix {
+    let (locator, source) = (checker.locator(), checker.source());
+    let comment_ranges = checker.comment_ranges();
+
+    let item_range = parenthesized_range(item.into(), iterable.into(), comment_ranges, source)
+        .unwrap_or(item.range());
+    let current_stmt_start = checker.semantic().current_statement().start();
+
+    let replace_op = Edit::range_replacement(
+        membership_test.replacement_op().to_string(),
+        membership_test.range(),
+    );
+    let mut other_edits: SmallVec<[Edit; 2]> = smallvec![];
+
+    let item_in_source = locator.slice(item_range);
+    let replace_iterable_with_item =
+        Edit::range_replacement(item_in_source.to_string(), iterable.range());
+
+    other_edits.push(replace_iterable_with_item);
+
+    let aggregated_comments =
+        merge_to_be_removed_comments(iterable, item_range, current_stmt_start, checker);
+
+    if !aggregated_comments.is_empty() {
+        let move_comments = Edit::insertion(aggregated_comments, current_stmt_start);
+        other_edits.push(move_comments);
+    }
+
+    let applicability = if iterable.is_string_literal_expr() {
+        Applicability::Unsafe
+    } else {
+        Applicability::Safe
+    };
+
+    Fix::applicable_edits(replace_op, other_edits, applicability)
+}
+
+fn merge_to_be_removed_comments(
+    iterable: &Expr,
+    item_range: TextRange,
+    current_stmt_start: TextSize,
+    checker: &Checker,
+) -> String {
+    let (locator, tokens) = (checker.locator(), checker.tokens());
+
+    let mut aggregated_comments = String::new();
+
+    let stmt_indentation_range =
+        TextRange::new(locator.line_start(current_stmt_start), current_stmt_start);
+    let stmt_indentation = locator.slice(stmt_indentation_range);
+    let line_ending = checker.stylist().line_ending().to_string();
+
+    let iterable_start_to_item_start = TextRange::new(iterable.start(), item_range.start());
+    let item_end_to_iterable_end = TextRange::new(item_range.end(), iterable.end());
+
+    tokens
+        .in_range(iterable_start_to_item_start)
+        .iter()
+        .chain(tokens.in_range(item_end_to_iterable_end))
+        .filter(|token| matches!(token.kind(), TokenKind::Comment))
+        .map(|token| locator.slice(token))
+        .for_each(|comment| {
+            aggregated_comments.push_str(&format!("{comment}{line_ending}{stmt_indentation}"));
+        });
+
+    aggregated_comments
+}
+
 #[derive(Debug, Clone, Copy, PartialEq, Eq)]
 enum MembershipTest {
     /// Ex) `1 in [1]`
-    In,
+    In { range: TextRange },
     /// Ex) `1 not in [1]`
-    NotIn,
+    NotIn { range: TextRange },
 }
 
 impl MembershipTest {
     /// Returns the replacement comparison operator for this membership test.
     fn replacement_op(self) -> CmpOp {
         match self {
-            MembershipTest::In => CmpOp::Eq,
-            MembershipTest::NotIn => CmpOp::NotEq,
+            Self::In { .. } => CmpOp::Eq,
+            Self::NotIn { .. } => CmpOp::NotEq,
+        }
+    }
+}
+
+impl Ranged for MembershipTest {
+    /// The original range of the operator
+    fn range(&self) -> TextRange {
+        match self {
+            Self::In { range } => *range,
+            Self::NotIn { range } => *range,
         }
     }
 }