[flake8-simplify] Add fix for if-with-same-arms (SIM114)

crates/ruff_linter/resources/test/fixtures/flake8_simplify/SIM114.py

@@ -4,6 +4,11 @@
 elif c:
     b
 
+if a:  # we preserve comments, too!
+    b
+elif c:  # but not on the second branch
+    b
+
 if x == 1:
     for _ in range(20):
         print("hello")
@@ -63,14 +68,18 @@
     errors = 1
 elif result.eofs == "E":
     errors = 1
+elif result.eofs == "X":
+    errors = 1
+elif result.eofs == "C":
+    errors = 1
 
 
 # OK
 def complicated_calc(*arg, **kwargs):
     return 42
 
 
-def foo(p):
+def func(p):
     if p == 2:
         return complicated_calc(microsecond=0)
     elif p == 3:
@@ -103,7 +112,7 @@ def foo(p):
     x = 1
 
 
-def foo():
+def func():
     a = True
     b = False
     if a > b:  # end-of-line
@@ -114,3 +123,23 @@ def foo():
         return 4
     elif b is None:
         return 4
+
+
+def func():
+    """Ensure that the named expression is parenthesized when merged."""
+    a = True
+    b = False
+    if a > b:  # end-of-line
+        return 3
+    elif a := 1:
+        return 3
+
+
+if a:  # we preserve comments, too!
+    b
+elif c:  # but not on the second branch
+    b
+
+
+if a: b  # here's a comment
+elif c: b

crates/ruff_linter/src/checkers/ast/analyze/statement.rs

@@ -1092,7 +1092,7 @@ pub(crate) fn statement(stmt: &Stmt, checker: &mut Checker) {
                 );
             }
             if checker.enabled(Rule::IfWithSameArms) {
-                flake8_simplify::rules::if_with_same_arms(checker, checker.locator, if_);
+                flake8_simplify::rules::if_with_same_arms(checker, if_);
             }
             if checker.enabled(Rule::NeedlessBool) {
                 flake8_simplify::rules::needless_bool(checker, if_);

crates/ruff_linter/src/rules/flake8_simplify/mod.rs

@@ -60,6 +60,7 @@ mod tests {
     #[test_case(Rule::YodaConditions, Path::new("SIM300.py"))]
     #[test_case(Rule::IfElseBlockInsteadOfDictGet, Path::new("SIM401.py"))]
     #[test_case(Rule::DictGetWithNoneDefault, Path::new("SIM910.py"))]
+    #[test_case(Rule::IfWithSameArms, Path::new("SIM114.py"))]
     fn preview_rules(rule_code: Rule, path: &Path) -> Result<()> {
         let snapshot = format!(
             "preview__{}_{}",

crates/ruff_linter/src/rules/flake8_simplify/rules/if_with_same_arms.rs

@@ -1,8 +1,15 @@
-use ruff_diagnostics::{Diagnostic, Violation};
+use std::borrow::Cow;
+
+use anyhow::Result;
+
+use ruff_diagnostics::{Diagnostic, Edit, Fix, FixAvailability, Violation};
 use ruff_macros::{derive_message_formats, violation};
-use ruff_python_ast as ast;
 use ruff_python_ast::comparable::ComparableStmt;
+use ruff_python_ast::parenthesize::parenthesized_range;
 use ruff_python_ast::stmt_if::{if_elif_branches, IfElifBranch};
+use ruff_python_ast::{self as ast, Expr};
+use ruff_python_index::Indexer;
+use ruff_python_trivia::{SimpleTokenKind, SimpleTokenizer};
 use ruff_source_file::Locator;
 use ruff_text_size::{Ranged, TextRange};
 
@@ -32,14 +39,20 @@ use crate::checkers::ast::Checker;
 pub struct IfWithSameArms;
 
 impl Violation for IfWithSameArms {
+    const FIX_AVAILABILITY: FixAvailability = FixAvailability::Sometimes;
+
     #[derive_message_formats]
     fn message(&self) -> String {
         format!("Combine `if` branches using logical `or` operator")
     }
+
+    fn fix_title(&self) -> Option<String> {
+        Some("Combine `if` branches".to_string())
+    }
 }
 
 /// SIM114
-pub(crate) fn if_with_same_arms(checker: &mut Checker, locator: &Locator, stmt_if: &ast::StmtIf) {
+pub(crate) fn if_with_same_arms(checker: &mut Checker, stmt_if: &ast::StmtIf) {
     let mut branches_iter = if_elif_branches(stmt_if).peekable();
     while let Some(current_branch) = branches_iter.next() {
         let Some(following_branch) = branches_iter.peek() else {
@@ -63,26 +76,101 @@ pub(crate) fn if_with_same_arms(checker: &mut Checker, locator: &Locator, stmt_i
         let first_comments = checker
             .indexer()
             .comment_ranges()
-            .comments_in_range(body_range(&current_branch, locator))
+            .comments_in_range(body_range(&current_branch, checker.locator()))
             .iter()
-            .map(|range| locator.slice(*range));
+            .map(|range| checker.locator().slice(*range));
         let second_comments = checker
             .indexer()
             .comment_ranges()
-            .comments_in_range(body_range(following_branch, locator))
+            .comments_in_range(body_range(following_branch, checker.locator()))
             .iter()
-            .map(|range| locator.slice(*range));
+            .map(|range| checker.locator().slice(*range));
         if !first_comments.eq(second_comments) {
             continue;
         }
 
-        checker.diagnostics.push(Diagnostic::new(
+        let mut diagnostic = Diagnostic::new(
             IfWithSameArms,
             TextRange::new(current_branch.start(), following_branch.end()),
-        ));
+        );
+
+        if checker.settings.preview.is_enabled() {
+            diagnostic.try_set_fix(|| {
+                merge_branches(
+                    stmt_if,
+                    &current_branch,
+                    following_branch,
+                    checker.locator(),
+                    checker.indexer(),
+                )
+            });
+        }
+
+        checker.diagnostics.push(diagnostic);
     }
 }
 
+/// Generate a [`Fix`] to merge two [`IfElifBranch`] branches.
+fn merge_branches(
+    stmt_if: &ast::StmtIf,
+    current_branch: &IfElifBranch,
+    following_branch: &IfElifBranch,
+    locator: &Locator,
+    indexer: &Indexer,
+) -> Result<Fix> {
+    // Identify the colon (`:`) at the end of the current branch's test.
+    let Some(current_branch_colon) =
+        SimpleTokenizer::starts_at(current_branch.test.end(), locator.contents())
+            .find(|token| token.kind == SimpleTokenKind::Colon)
+    else {
+        return Err(anyhow::anyhow!("Expected colon after test"));
+    };
+
+    let mut following_branch_tokenizer =
+        SimpleTokenizer::starts_at(following_branch.test.end(), locator.contents());
+
+    // Identify the colon (`:`) at the end of the following branch's test.
+    let Some(following_branch_colon) =
+        following_branch_tokenizer.find(|token| token.kind == SimpleTokenKind::Colon)
+    else {
+        return Err(anyhow::anyhow!("Expected colon after test"));
+    };
+
+    let main_edit = Edit::deletion(
+        locator.full_line_end(current_branch_colon.end()),
+        locator.full_line_end(following_branch_colon.end()),
+    );
+
+    // If the test isn't parenthesized, consider parenthesizing it.
+    let following_branch_test = if let Some(range) = parenthesized_range(
+        following_branch.test.into(),
+        stmt_if.into(),
+        indexer.comment_ranges(),
+        locator.contents(),
+    ) {
+        Cow::Borrowed(locator.slice(range))
+    } else if matches!(
+        following_branch.test,
+        Expr::BoolOp(ast::ExprBoolOp {
+            op: ast::BoolOp::Or,
+            ..
+        }) | Expr::Lambda(_)
+            | Expr::NamedExpr(_)
+    ) {
+        Cow::Owned(format!("({})", locator.slice(following_branch.test)))
+    } else {
+        Cow::Borrowed(locator.slice(following_branch.test))
+    };
+
+    Ok(Fix::safe_edits(
+        main_edit,
+        [Edit::insertion(
+            format!(" or {following_branch_test}"),
+            current_branch_colon.start(),
+        )],
+    ))
+}
+
 /// Return the [`TextRange`] of an [`IfElifBranch`]'s body (from the end of the test to the end of
 /// the body).
 fn body_range(branch: &IfElifBranch, locator: &Locator) -> TextRange {