[WIP] Editing around comments

crates/ruff_linter/src/fix/edits.rs

@@ -2,6 +2,7 @@
 
 use anyhow::{Context, Result};
 
+use itertools::Itertools;
 use ruff_diagnostics::Edit;
 use ruff_python_ast::parenthesize::parenthesized_range;
 use ruff_python_ast::{self as ast, Arguments, ExceptHandler, Expr, ExprList, Parameters, Stmt};
@@ -10,8 +11,8 @@ use ruff_python_codegen::Stylist;
 use ruff_python_index::Indexer;
 use ruff_python_trivia::textwrap::dedent_to;
 use ruff_python_trivia::{
-    has_leading_content, is_python_whitespace, CommentRanges, PythonWhitespace, SimpleTokenKind,
-    SimpleTokenizer,
+    has_leading_content, indentation_at_offset, is_python_whitespace, CommentRanges,
+    PythonWhitespace, SimpleTokenKind, SimpleTokenizer,
 };
 use ruff_source_file::{LineRanges, NewlineWithTrailingNewline, UniversalNewlines};
 use ruff_text_size::{Ranged, TextLen, TextRange, TextSize};
@@ -107,6 +108,82 @@ pub(crate) fn delete_comment(range: TextRange, locator: &Locator) -> Edit {
     }
 }
 
+pub(crate) fn delete_around_comments(
+    start: TextSize,
+    end: TextSize,
+    comment_ranges: &CommentRanges,
+    source: &str,
+) -> Vec<Edit> {
+    let delete_range = TextRange::new(start, end);
+    // Begin by restricting attention to the comment ranges
+    // that intersect the deletion range. Suppose it looks like:
+    //
+    // (s1,e1),(s2,e2),...,(sn,en)
+    //
+    // Flatten, then prepend and append by the start and
+    // end of the deletion range to get:
+    //
+    // s0, s1, e1, s2, ..., sn, en, e0
+    //
+    // Now pair up to get the desired deletion ranges
+    // in the complement of the comments:
+    //
+    // (s0,s1), (e1,s2), ..., (en,e0)
+    let mut edits: Vec<Edit> = std::iter::once(start)
+        .chain(
+            comment_ranges
+                .iter()
+                .filter(|range| range.intersect(delete_range).is_some())
+                // extend comment range to include last newline
+                .flat_map(|range| [range.start(), range.end() + TextSize::from(1)]),
+        )
+        .chain(std::iter::once(end))
+        // Here we clamp the comment intervals to the deletion range
+        // (somewhat awkwardly: each conditional branch triggers
+        // at most once, and the remaining is a no-op.)
+        .map(|locn| {
+            if locn < start {
+                return start;
+            }
+            if locn > end {
+                return end;
+            }
+            locn
+        })
+        .tuples::<(_, _)>()
+        .map(|(left, right)| Edit::deletion(left, right))
+        .collect();
+    // Adjust the last deletion so remaining content matches the
+    // indentation at start of original deletion range.
+    let Some(last_edit) = edits.last_mut() else {
+        return edits;
+    };
+    let Some(start_indent) = indentation_at_offset(start, source) else {
+        return edits;
+    };
+    if !start_indent.is_empty() {
+        *last_edit = Edit::range_replacement(start_indent.to_string(), last_edit.range());
+    };
+    edits
+}
+
+pub(crate) fn replace_around_comments(
+    content: &str,
+    start: TextSize,
+    end: TextSize,
+    comment_ranges: &CommentRanges,
+    source: &str,
+) -> Vec<Edit> {
+    let mut edits = delete_around_comments(start, end, comment_ranges, source);
+    let Some(last_edit) = edits.last_mut() else {
+        return edits;
+    };
+    *last_edit = Edit::range_replacement(
+        format!("{}{}", last_edit.content().unwrap_or_default(), content),
+        last_edit.range(),
+    );
+    edits
+}
 /// Generate a `Fix` to remove the specified imports from an `import` statement.
 pub(crate) fn remove_unused_imports<'a>(
     member_names: impl Iterator<Item = &'a str>,

crates/ruff_linter/src/rules/flake8_comprehensions/rules/unnecessary_literal_within_list_call.rs

@@ -1,3 +1,4 @@
+use crate::fix::edits::delete_around_comments;
 use ruff_diagnostics::{AlwaysFixableViolation, Diagnostic, Edit, Fix};
 use ruff_macros::{derive_message_formats, ViolationMetadata};
 use ruff_python_ast::{self as ast, Expr};
@@ -89,10 +90,20 @@ pub(crate) fn unnecessary_literal_within_list_call(checker: &mut Checker, call:
     // Convert `list([1, 2])` to `[1, 2]`
     let fix = {
         // Delete from the start of the call to the start of the argument.
-        let call_start = Edit::deletion(call.start(), argument.start());
+        let mut call_start_edits = delete_around_comments(
+            call.start(),
+            argument.start(),
+            checker.comment_ranges(),
+            checker.locator().contents(),
+        );
 
         // Delete from the end of the argument to the end of the call.
-        let call_end = Edit::deletion(argument.end(), call.end());
+        let call_end_edits = delete_around_comments(
+            argument.end(),
+            call.end(),
+            checker.comment_ranges(),
+            checker.locator().contents(),
+        );
 
         // If this is a tuple, we also need to convert the inner argument to a list.
         if argument.is_tuple_expr() {
@@ -109,10 +120,18 @@ pub(crate) fn unnecessary_literal_within_list_call(checker: &mut Checker, call:
                 argument.end() - TextSize::from(1),
                 argument.end(),
             );
-
-            Fix::unsafe_edits(call_start, [argument_start, argument_end, call_end])
+            call_start_edits.extend(
+                std::iter::once(argument_start)
+                    .chain(std::iter::once(argument_end))
+                    .chain(call_end_edits),
+            );
+            let (edit, rest) = call_start_edits.split_first().unwrap();
+            Fix::unsafe_edits(edit.clone(), rest.to_owned())
         } else {
-            Fix::unsafe_edits(call_start, [call_end])
+            call_start_edits.extend(call_end_edits);
+
+            let (edit, rest) = call_start_edits.split_first().unwrap();
+            Fix::unsafe_edits(edit.clone(), rest.to_owned())
         }
     };
 

crates/ruff_linter/src/rules/flake8_comprehensions/rules/unnecessary_literal_within_tuple_call.rs

@@ -1,11 +1,12 @@
-use ruff_diagnostics::{AlwaysFixableViolation, Diagnostic, Edit, Fix};
+use ruff_diagnostics::{AlwaysFixableViolation, Diagnostic, Fix};
 use ruff_macros::{derive_message_formats, ViolationMetadata};
 use ruff_python_ast::helpers::any_over_expr;
 use ruff_python_ast::{self as ast, Expr};
 use ruff_python_trivia::{SimpleTokenKind, SimpleTokenizer};
 use ruff_text_size::{Ranged, TextRange, TextSize};
 
 use crate::checkers::ast::Checker;
+use crate::fix::edits::replace_around_comments;
 use crate::rules::flake8_comprehensions::fixes;
 
 use super::helpers;
@@ -129,22 +130,24 @@ pub(crate) fn unnecessary_literal_within_tuple_call(
                 };
 
                 // Replace `[` with `(`.
-                let elt_start = Edit::replacement(
-                    "(".into(),
+                let mut elt_start = replace_around_comments(
+                    "(",
                     call.start(),
                     argument.start() + TextSize::from(1),
+                    checker.comment_ranges(),
+                    checker.locator().contents(),
                 );
                 // Replace `]` with `)` or `,)`.
-                let elt_end = Edit::replacement(
-                    if needs_trailing_comma {
-                        ",)".into()
-                    } else {
-                        ")".into()
-                    },
+                let elt_end = replace_around_comments(
+                    if needs_trailing_comma { ",)" } else { ")" },
                     argument.end() - TextSize::from(1),
                     call.end(),
+                    checker.comment_ranges(),
+                    checker.locator().contents(),
                 );
-                Fix::unsafe_edits(elt_start, [elt_end])
+                elt_start.extend(elt_end);
+                let (edit, rest) = elt_start.split_first().unwrap();
+                Fix::unsafe_edits(edit.clone(), rest.to_owned())
             });
         }
 

crates/ruff_linter/src/rules/flake8_comprehensions/snapshots/ruff_linter__rules__flake8_comprehensions__tests__C409_C409.py.snap

@@ -130,10 +130,11 @@ C409.py:12:1: C409 [*] Unnecessary list literal passed to `tuple()` (rewrite as
 12    |-tuple(  # comment
 13    |-    [1, 2]
 14    |-)
-   12 |+(1, 2)
-15 13 | 
-16 14 | tuple([  # comment
-17 15 |     1, 2
+   12 |+# comment
+   13 |+(1, 2)
+15 14 | 
+16 15 | tuple([  # comment
+17 16 |     1, 2
 
 C409.py:16:1: C409 [*] Unnecessary list literal passed to `tuple()` (rewrite as a tuple literal)
    |

crates/ruff_linter/src/rules/flake8_comprehensions/snapshots/ruff_linter__rules__flake8_comprehensions__tests__C410_C410.py.snap

@@ -91,10 +91,11 @@ C410.py:7:1: C410 [*] Unnecessary list literal passed to `list()` (remove the ou
 7     |-list(  # comment
 8     |-    [1, 2]
 9     |-)
-   7  |+[1, 2]
-10 8  | 
-11 9  | list([  # comment
-12 10 |     1, 2
+   7  |+# comment
+   8  |+[1, 2]
+10 9  | 
+11 10 | list([  # comment
+12 11 |     1, 2
 
 C410.py:11:1: C410 [*] Unnecessary list literal passed to `list()` (remove the outer call to `list()`)
    |

crates/ruff_linter/src/rules/flake8_comprehensions/snapshots/ruff_linter__rules__flake8_comprehensions__tests__preview__C409_C409.py.snap

@@ -130,10 +130,11 @@ C409.py:12:1: C409 [*] Unnecessary list literal passed to `tuple()` (rewrite as
 12    |-tuple(  # comment
 13    |-    [1, 2]
 14    |-)
-   12 |+(1, 2)
-15 13 | 
-16 14 | tuple([  # comment
-17 15 |     1, 2
+   12 |+# comment
+   13 |+(1, 2)
+15 14 | 
+16 15 | tuple([  # comment
+17 16 |     1, 2
 
 C409.py:16:1: C409 [*] Unnecessary list literal passed to `tuple()` (rewrite as a tuple literal)
    |

crates/ruff_linter/src/rules/pyupgrade/rules/convert_typed_dict_functional_to_class.rs

@@ -1,4 +1,4 @@
-use ruff_diagnostics::{Applicability, Diagnostic, Edit, Fix, FixAvailability, Violation};
+use ruff_diagnostics::{Diagnostic, Fix, FixAvailability, Violation};
 use ruff_macros::{derive_message_formats, ViolationMetadata};
 use ruff_python_ast::helpers::is_dunder;
 use ruff_python_ast::{self as ast, Arguments, Expr, ExprContext, Identifier, Keyword, Stmt};
@@ -10,6 +10,7 @@ use ruff_source_file::LineRanges;
 use ruff_text_size::{Ranged, TextRange};
 
 use crate::checkers::ast::Checker;
+use crate::fix::edits::replace_around_comments;
 
 /// ## What it does
 /// Checks for `TypedDict` declarations that use functional syntax.
@@ -98,6 +99,7 @@ pub(crate) fn convert_typed_dict_functional_to_class(
             base_class,
             checker.generator(),
             checker.comment_ranges(),
+            checker.locator().contents(),
         ));
     }
     checker.diagnostics.push(diagnostic);
@@ -265,6 +267,7 @@ fn match_fields_and_total(arguments: &Arguments) -> Option<(Vec<Stmt>, Option<&K
 }
 
 /// Generate a `Fix` to convert a `TypedDict` from functional to class.
+#[allow(clippy::too_many_arguments)]
 fn convert_to_class(
     stmt: &Stmt,
     class_name: &str,
@@ -273,21 +276,22 @@ fn convert_to_class(
     base_class: &Expr,
     generator: Generator,
     comment_ranges: &CommentRanges,
+    source: &str,
 ) -> Fix {
-    Fix::applicable_edit(
-        Edit::range_replacement(
-            generator.stmt(&create_class_def_stmt(
-                class_name,
-                body,
-                total_keyword,
-                base_class,
-            )),
-            stmt.range(),
-        ),
-        if comment_ranges.intersects(stmt.range()) {
-            Applicability::Unsafe
-        } else {
-            Applicability::Safe
-        },
-    )
+    let edits = replace_around_comments(
+        &generator.stmt(&create_class_def_stmt(
+            class_name,
+            body,
+            total_keyword,
+            base_class,
+        )),
+        stmt.start(),
+        stmt.end(),
+        comment_ranges,
+        source,
+    );
+    let Some((edit, rest)) = edits.split_first() else {
+        panic!()
+    };
+    Fix::safe_edits(edit.clone(), rest.to_owned())
 }

crates/ruff_linter/src/rules/pyupgrade/snapshots/ruff_linter__rules__pyupgrade__tests__UP013.py.snap

@@ -267,12 +267,13 @@ UP013.py:46:1: UP013 [*] Convert `X` from `TypedDict` functional to class syntax
    |
    = help: Convert `X` to class syntax
 
-ℹ Unsafe fix
+ℹ Safe fix
 43 43 | MyType = TypedDict("MyType", dict())
 44 44 | 
 45 45 | # Unsafe fix if comments are present
 46    |-X = TypedDict("X", {
 47    |-    "some_config": int,  # important
 48    |-})
-   46 |+class X(TypedDict):
-   47 |+    some_config: int
+   46 |+# important
+   47 |+class X(TypedDict):
+   48 |+    some_config: int