Allow diagnostics to generate multi-edit fixes

crates/ruff/src/autofix/mod.rs

@@ -4,7 +4,7 @@ use itertools::Itertools;
 use rustc_hash::FxHashMap;
 use rustpython_parser::ast::Location;
 
-use ruff_diagnostics::{Diagnostic, Edit};
+use ruff_diagnostics::{Diagnostic, Edit, Fix};
 use ruff_python_ast::source_code::Locator;
 use ruff_python_ast::types::Range;
 
@@ -15,7 +15,7 @@ pub mod helpers;
 
 /// Auto-fix errors in a file, and write the fixed source code to disk.
 pub fn fix_file(diagnostics: &[Diagnostic], locator: &Locator) -> Option<(String, FixTable)> {
-    if diagnostics.iter().all(|check| check.fix.is_none()) {
+    if diagnostics.iter().all(|check| check.fix.is_empty()) {
         None
     } else {
         Some(apply_fixes(diagnostics.iter(), locator))
@@ -34,36 +34,43 @@ fn apply_fixes<'a>(
 
     for (rule, fix) in diagnostics
         .filter_map(|diagnostic| {
-            diagnostic
-                .fix
-                .as_ref()
-                .map(|fix| (diagnostic.kind.rule(), fix))
+            if diagnostic.fix.is_empty() {
+                None
+            } else {
+                Some((diagnostic.kind.rule(), &diagnostic.fix))
+            }
         })
         .sorted_by(|(rule1, fix1), (rule2, fix2)| cmp_fix(*rule1, *rule2, fix1, fix2))
     {
         // If we already applied an identical fix as part of another correction, skip
         // any re-application.
-        if applied.contains(&fix) {
+        if fix.edits().iter().all(|edit| applied.contains(edit)) {
             *fixed.entry(rule).or_default() += 1;
             continue;
         }
 
         // Best-effort approach: if this fix overlaps with a fix we've already applied,
         // skip it.
-        if last_pos.map_or(false, |last_pos| last_pos >= fix.location) {
+        if last_pos.map_or(false, |last_pos| {
+            fix.location()
+                .map_or(false, |fix_location| last_pos >= fix_location)
+        }) {
             continue;
         }
 
-        // Add all contents from `last_pos` to `fix.location`.
-        let slice = locator.slice(Range::new(last_pos.unwrap_or_default(), fix.location));
-        output.push_str(slice);
+        for edit in fix.edits() {
+            // Add all contents from `last_pos` to `fix.location`.
+            let slice = locator.slice(Range::new(last_pos.unwrap_or_default(), edit.location));
+            output.push_str(slice);
 
-        // Add the patch itself.
-        output.push_str(&fix.content);
+            // Add the patch itself.
+            output.push_str(&edit.content);
+
+            // Track that the edit was applied.
+            last_pos = Some(edit.end_location);
+            applied.insert(edit);
+        }
 
-        // Track that the fix was applied.
-        last_pos = Some(fix.end_location);
-        applied.insert(fix);
         *fixed.entry(rule).or_default() += 1;
     }
 
@@ -93,9 +100,9 @@ pub(crate) fn apply_fix(fix: &Edit, locator: &Locator) -> String {
 }
 
 /// Compare two fixes.
-fn cmp_fix(rule1: Rule, rule2: Rule, fix1: &Edit, fix2: &Edit) -> std::cmp::Ordering {
-    fix1.location
-        .cmp(&fix2.location)
+fn cmp_fix(rule1: Rule, rule2: Rule, fix1: &Fix, fix2: &Fix) -> std::cmp::Ordering {
+    fix1.location()
+        .cmp(&fix2.location())
         .then_with(|| match (&rule1, &rule2) {
             // Apply `EndsInPeriod` fixes before `NewLineAfterLastParagraph` fixes.
             (Rule::EndsInPeriod, Rule::NewLineAfterLastParagraph) => std::cmp::Ordering::Less,
@@ -115,15 +122,14 @@ mod tests {
     use crate::autofix::{apply_fix, apply_fixes};
     use crate::rules::pycodestyle::rules::MissingNewlineAtEndOfFile;
 
-    fn create_diagnostics(fixes: impl IntoIterator<Item = Edit>) -> Vec<Diagnostic> {
-        fixes
-            .into_iter()
-            .map(|fix| Diagnostic {
+    fn create_diagnostics(edit: impl IntoIterator<Item = Edit>) -> Vec<Diagnostic> {
+        edit.into_iter()
+            .map(|edit| Diagnostic {
                 // The choice of rule here is arbitrary.
                 kind: MissingNewlineAtEndOfFile.into(),
-                location: fix.location,
-                end_location: fix.end_location,
-                fix: Some(fix),
+                location: edit.location,
+                end_location: edit.end_location,
+                fix: edit.into(),
                 parent: None,
             })
             .collect()

crates/ruff/src/checkers/ast/mod.rs

@@ -2455,7 +2455,7 @@ where
                     pyupgrade::rules::replace_universal_newlines(self, func, keywords);
                 }
                 if self.settings.rules.enabled(Rule::ReplaceStdoutStderr) {
-                    pyupgrade::rules::replace_stdout_stderr(self, expr, func, keywords);
+                    pyupgrade::rules::replace_stdout_stderr(self, expr, func, args, keywords);
                 }
                 if self.settings.rules.enabled(Rule::OSErrorAlias) {
                     pyupgrade::rules::os_error_alias_call(self, func);

crates/ruff/src/checkers/logical_lines.rs

@@ -5,7 +5,7 @@ use itertools::Itertools;
 use rustpython_parser::ast::Location;
 use rustpython_parser::lexer::LexResult;
 
-use ruff_diagnostics::Diagnostic;
+use ruff_diagnostics::{Diagnostic, Fix};
 use ruff_python_ast::source_code::{Locator, Stylist};
 use ruff_python_ast::types::Range;
 
@@ -75,7 +75,7 @@ pub fn check_logical_lines(
                         kind,
                         location,
                         end_location: location,
-                        fix: None,
+                        fix: Fix::empty(),
                         parent: None,
                     });
                 }
@@ -93,7 +93,7 @@ pub fn check_logical_lines(
                         kind,
                         location,
                         end_location: location,
-                        fix: None,
+                        fix: Fix::empty(),
                         parent: None,
                     });
                 }
@@ -108,7 +108,7 @@ pub fn check_logical_lines(
                         kind,
                         location,
                         end_location: location,
-                        fix: None,
+                        fix: Fix::empty(),
                         parent: None,
                     });
                 }
@@ -120,7 +120,7 @@ pub fn check_logical_lines(
                         kind,
                         location,
                         end_location: location,
-                        fix: None,
+                        fix: Fix::empty(),
                         parent: None,
                     });
                 }
@@ -133,7 +133,7 @@ pub fn check_logical_lines(
                         kind,
                         location: range.location,
                         end_location: range.end_location,
-                        fix: None,
+                        fix: Fix::empty(),
                         parent: None,
                     });
                 }
@@ -148,7 +148,7 @@ pub fn check_logical_lines(
                         kind,
                         location,
                         end_location: location,
-                        fix: None,
+                        fix: Fix::empty(),
                         parent: None,
                     });
                 }
@@ -159,7 +159,7 @@ pub fn check_logical_lines(
                         kind,
                         location,
                         end_location: location,
-                        fix: None,
+                        fix: Fix::empty(),
                         parent: None,
                     });
                 }
@@ -210,7 +210,7 @@ pub fn check_logical_lines(
                     kind,
                     location,
                     end_location: location,
-                    fix: None,
+                    fix: Fix::empty(),
                     parent: None,
                 });
             }

crates/ruff/src/message.rs

@@ -3,7 +3,7 @@ use std::cmp::Ordering;
 pub use rustpython_parser::ast::Location;
 use serde::{Deserialize, Serialize};
 
-use ruff_diagnostics::{Diagnostic, DiagnosticKind, Edit};
+use ruff_diagnostics::{Diagnostic, DiagnosticKind, Fix};
 use ruff_python_ast::source_code::Locator;
 use ruff_python_ast::types::Range;
 
@@ -12,7 +12,7 @@ pub struct Message {
     pub kind: DiagnosticKind,
     pub location: Location,
     pub end_location: Location,
-    pub fix: Option<Edit>,
+    pub fix: Fix,
     pub filename: String,
     pub source: Option<Source>,
     pub noqa_row: usize,

crates/ruff/src/rules/pyupgrade/rules/replace_stdout_stderr.rs

@@ -1,12 +1,13 @@
+use anyhow::Result;
 use rustpython_parser::ast::{Expr, Keyword};
 
-use ruff_diagnostics::{AlwaysAutofixableViolation, Diagnostic, Edit};
+use ruff_diagnostics::{AlwaysAutofixableViolation, Diagnostic, Edit, Fix};
 use ruff_macros::{derive_message_formats, violation};
 use ruff_python_ast::helpers::find_keyword;
-use ruff_python_ast::source_code::{Locator, Stylist};
+use ruff_python_ast::source_code::Locator;
 use ruff_python_ast::types::Range;
-use ruff_python_ast::whitespace::indentation;
 
+use crate::autofix::helpers::remove_argument;
 use crate::checkers::ast::Checker;
 use crate::registry::AsRule;
 
@@ -24,87 +25,46 @@ impl AlwaysAutofixableViolation for ReplaceStdoutStderr {
     }
 }
 
-#[derive(Debug)]
-struct MiddleContent<'a> {
-    contents: &'a str,
-    multi_line: bool,
-}
-
-/// Return the number of "dirty" characters.
-fn dirty_count(iter: impl Iterator<Item = char>) -> usize {
-    let mut the_count = 0;
-    for current_char in iter {
-        if current_char == ' '
-            || current_char == ','
-            || current_char == '\n'
-            || current_char == '\r'
-        {
-            the_count += 1;
-        } else {
-            break;
-        }
-    }
-    the_count
-}
-
-/// Extract the `Middle` content between two arguments.
-fn extract_middle(contents: &str) -> Option<MiddleContent> {
-    let multi_line = contents.contains('\n');
-    let start_gap = dirty_count(contents.chars());
-    if contents.len() == start_gap {
-        return None;
-    }
-    let end_gap = dirty_count(contents.chars().rev());
-    Some(MiddleContent {
-        contents: &contents[start_gap..contents.len() - end_gap],
-        multi_line,
-    })
-}
-
 /// Generate a [`Edit`] for a `stdout` and `stderr` [`Keyword`] pair.
 fn generate_fix(
-    stylist: &Stylist,
     locator: &Locator,
+    func: &Expr,
+    args: &[Expr],
+    keywords: &[Keyword],
     stdout: &Keyword,
     stderr: &Keyword,
-) -> Option<Edit> {
-    let line_end = stylist.line_ending().as_str();
-    let first = if stdout.location < stderr.location {
-        stdout
-    } else {
-        stderr
-    };
-    let last = if stdout.location > stderr.location {
-        stdout
+) -> Result<Fix> {
+    let (first, second) = if stdout.location < stderr.location {
+        (stdout, stderr)
     } else {
-        stderr
+        (stderr, stdout)
     };
-    let mut contents = String::from("capture_output=True");
-    if let Some(middle) =
-        extract_middle(locator.slice(Range::new(first.end_location.unwrap(), last.location)))
-    {
-        if middle.multi_line {
-            let Some(indent) = indentation(locator, first) else {
-                return None;
-            };
-            contents.push(',');
-            contents.push_str(line_end);
-            contents.push_str(indent);
-        } else {
-            contents.push(',');
-            contents.push(' ');
-        }
-        contents.push_str(middle.contents);
-    }
-    Some(Edit::replacement(
-        contents,
-        first.location,
-        last.end_location.unwrap(),
-    ))
+    Ok(Fix::new(vec![
+        Edit::replacement(
+            "capture_output=True".to_string(),
+            first.location,
+            first.end_location.unwrap(),
+        ),
+        remove_argument(
+            locator,
+            func.location,
+            second.location,
+            second.end_location.unwrap(),
+            args,
+            keywords,
+            false,
+        )?,
+    ]))
 }
 
 /// UP022
-pub fn replace_stdout_stderr(checker: &mut Checker, expr: &Expr, func: &Expr, kwargs: &[Keyword]) {
+pub fn replace_stdout_stderr(
+    checker: &mut Checker,
+    expr: &Expr,
+    func: &Expr,
+    args: &[Expr],
+    keywords: &[Keyword],
+) {
     if checker
         .ctx
         .resolve_call_path(func)
@@ -113,10 +73,10 @@ pub fn replace_stdout_stderr(checker: &mut Checker, expr: &Expr, func: &Expr, kw
         })
     {
         // Find `stdout` and `stderr` kwargs.
-        let Some(stdout) = find_keyword(kwargs, "stdout") else {
+        let Some(stdout) = find_keyword(keywords, "stdout") else {
             return;
         };
-        let Some(stderr) = find_keyword(kwargs, "stderr") else {
+        let Some(stderr) = find_keyword(keywords, "stderr") else {
             return;
         };
 
@@ -139,9 +99,9 @@ pub fn replace_stdout_stderr(checker: &mut Checker, expr: &Expr, func: &Expr, kw
 
         let mut diagnostic = Diagnostic::new(ReplaceStdoutStderr, Range::from(expr));
         if checker.patch(diagnostic.kind.rule()) {
-            if let Some(fix) = generate_fix(checker.stylist, checker.locator, stdout, stderr) {
-                diagnostic.set_fix(fix);
-            };
+            diagnostic.try_set_fix(|| {
+                generate_fix(checker.locator, func, args, keywords, stdout, stderr)
+            });
         }
         checker.diagnostics.push(diagnostic);
     }

crates/ruff/src/test.rs

@@ -51,7 +51,7 @@ pub fn test_path(path: impl AsRef<Path>, settings: &Settings) -> Result<Vec<Diag
     // Detect autofixes that don't converge after multiple iterations.
     if diagnostics
         .iter()
-        .any(|diagnostic| diagnostic.fix.is_some())
+        .any(|diagnostic| !diagnostic.fix.is_empty())
     {
         let max_iterations = 10;