Autofix SIM102 (NestedIfStatements)

Co-authored-by: Charlie Marsh <charlie.r.marsh@gmail.com>
Signed-off-by: Anders Kaseorg <andersk@mit.edu>

README.md

@@ -1004,7 +1004,7 @@ For more, see [flake8-simplify](https://pypi.org/project/flake8-simplify/0.19.3/
 | Code | Name | Message | Fix |
 | ---- | ---- | ------- | --- |
 | SIM101 | DuplicateIsinstanceCall | Multiple `isinstance` calls for `...`, merge into a single call | 🛠 |
-| SIM102 | NestedIfStatements | Use a single `if` statement instead of nested `if` statements |  |
+| SIM102 | NestedIfStatements | Use a single `if` statement instead of nested `if` statements | 🛠 |
 | SIM103 | ReturnBoolConditionDirectly | Return the condition `...` directly | 🛠 |
 | SIM105 | UseContextlibSuppress | Use `contextlib.suppress(...)` instead of try-except-pass |  |
 | SIM107 | ReturnInTryExceptFinally | Don't use `return` in `try`/`except` and `finally` |  |

resources/test/fixtures/flake8_simplify/SIM102.py

@@ -1,25 +1,88 @@
-if a:  # SIM102
+# SIM102
+if a:
     if b:
         c
 
-
+# SIM102
 if a:
     pass
-elif b:  # SIM102
+elif b:
     if c:
         d
 
+# SIM102
+if a:
+    # Unfixable due to placement of this comment.
+    if b:
+        c
+
+# SIM102
+if a:
+    if b:
+        # Fixable due to placement of this comment.
+        c
+
+# OK
 if a:
     if b:
         c
     else:
         d
 
+# OK
 if __name__ == "__main__":
     if foo():
         ...
 
+# OK
 if a:
     d
     if b:
         c
+
+while True:
+    # SIM102
+    if True:
+        if True:
+            """this
+is valid"""
+
+            """the indentation on
+            this line is significant"""
+
+            "this is" \
+"allowed too"
+
+            ("so is"
+"this for some reason")
+
+
+# SIM102
+if True:
+    if True:
+        """this
+is valid"""
+
+        """the indentation on
+        this line is significant"""
+
+        "this is" \
+"allowed too"
+
+        ("so is"
+"this for some reason")
+
+while True:
+    # SIM102
+    if node.module:
+        if node.module == "multiprocessing" or node.module.startswith(
+            "multiprocessing."
+        ):
+            print("Bad module!")
+
+# SIM102
+if node.module:
+    if node.module == "multiprocessing" or node.module.startswith(
+        "multiprocessing."
+    ):
+        print("Bad module!")

src/rules/flake8_simplify/rules/ast_if.rs

@@ -1,3 +1,4 @@
+use log::error;
 use rustpython_ast::{Cmpop, Constant, Expr, ExprContext, ExprKind, Stmt, StmtKind};
 
 use crate::ast::comparable::ComparableExpr;
@@ -9,6 +10,7 @@ use crate::ast::types::Range;
 use crate::checkers::ast::Checker;
 use crate::fix::Fix;
 use crate::registry::{Diagnostic, RuleCode};
+use crate::rules::flake8_simplify::rules::fix_if;
 use crate::violations;
 
 fn is_main_check(expr: &Expr) -> bool {
@@ -64,10 +66,30 @@ pub fn nested_if_statements(checker: &mut Checker, stmt: &Stmt) {
         return;
     }
 
-    checker.diagnostics.push(Diagnostic::new(
+    let mut diagnostic = Diagnostic::new(violations::NestedIfStatements, Range::from_located(stmt));
-        violations::NestedIfStatements,
+    if checker.patch(&RuleCode::SIM102) {
-        Range::from_located(stmt),
+        // The fixer preserves comments in the nested body, but removes comments between
-    ));
+        // the outer and inner if statements.
+        let nested_if = &body[0];
+        if !has_comments_in(
+            Range::new(stmt.location, nested_if.location),
+            checker.locator,
+        ) {
+            match fix_if::fix_nested_if_statements(checker.locator, stmt) {
+                Ok(fix) => {
+                    if fix
+                        .content
+                        .lines()
+                        .all(|line| line.len() <= checker.settings.line_length)
+                    {
+                        diagnostic.amend(fix);
+                    }
+                }
+                Err(err) => error!("Failed to fix nested if: {err}"),
+            }
+        }
+    }
+    checker.diagnostics.push(diagnostic);
 }
 
 fn is_one_line_return_bool(stmts: &[Stmt]) -> bool {

src/rules/flake8_simplify/rules/fix_if.rs

@@ -0,0 +1,137 @@
+use std::borrow::Cow;
+
+use anyhow::{bail, Result};
+use libcst_native::{
+    BooleanOp, BooleanOperation, Codegen, CodegenState, CompoundStatement, Expression, If,
+    LeftParen, ParenthesizableWhitespace, ParenthesizedNode, RightParen, SimpleWhitespace,
+    Statement, Suite,
+};
+use rustpython_ast::Location;
+
+use crate::ast::types::Range;
+use crate::ast::whitespace;
+use crate::cst::matchers::match_module;
+use crate::fix::Fix;
+use crate::source_code::Locator;
+
+fn parenthesize_and_operand(expr: Expression) -> Expression {
+    match &expr {
+        _ if !expr.lpar().is_empty() => expr,
+        Expression::BooleanOperation(boolean_operation)
+            if matches!(boolean_operation.operator, BooleanOp::Or { .. }) =>
+        {
+            expr.with_parens(LeftParen::default(), RightParen::default())
+        }
+        Expression::IfExp(_) | Expression::Lambda(_) | Expression::NamedExpr(_) => {
+            expr.with_parens(LeftParen::default(), RightParen::default())
+        }
+        _ => expr,
+    }
+}
+
+/// (SIM102) Convert `if a: if b:` to `if a and b:`.
+pub(crate) fn fix_nested_if_statements(
+    locator: &Locator,
+    stmt: &rustpython_ast::Stmt,
+) -> Result<Fix> {
+    // Infer the indentation of the outer block.
+    let Some(outer_indent) = whitespace::indentation(locator, stmt) else {
+        bail!("Unable to fix multiline statement");
+    };
+
+    // Extract the module text.
+    let contents = locator.slice_source_code_range(&Range::new(
+        Location::new(stmt.location.row(), 0),
+        Location::new(stmt.end_location.unwrap().row() + 1, 0),
+    ));
+
+    // Handle `elif` blocks differently; detect them upfront.
+    let is_elif = contents.trim_start().starts_with("elif");
+
+    // If this is an `elif`, we have to remove the `elif` keyword for now. (We'll
+    // restore the `el` later on.)
+    let module_text = if is_elif {
+        Cow::Owned(contents.replacen("elif", "if", 1))
+    } else {
+        contents
+    };
+
+    // If the block is indented, "embed" it in a function definition, to preserve
+    // indentation while retaining valid source code. (We'll strip the prefix later
+    // on.)
+    let module_text = if outer_indent.is_empty() {
+        module_text
+    } else {
+        Cow::Owned(format!("def f():\n{module_text}"))
+    };
+
+    // Parse the CST.
+    let mut tree = match_module(&module_text)?;
+
+    let statements = if outer_indent.is_empty() {
+        &mut *tree.body
+    } else {
+        let [Statement::Compound(CompoundStatement::FunctionDef(embedding))] = &mut *tree.body else {
+            bail!("Expected statement to be embedded in a function definition")
+        };
+
+        let Suite::IndentedBlock(indented_block) = &mut embedding.body else {
+            bail!("Expected indented block")
+        };
+        indented_block.indent = Some(&outer_indent);
+
+        &mut *indented_block.body
+    };
+
+    let [Statement::Compound(CompoundStatement::If(outer_if))] = statements else {
+        bail!("Expected one outer if statement")
+    };
+
+    let If {
+        body: Suite::IndentedBlock(ref mut outer_body),
+        orelse: None,
+        ..
+    } = outer_if else {
+        bail!("Expected outer if to have indented body and no else")
+    };
+
+    let [Statement::Compound(CompoundStatement::If(inner_if @ If { orelse: None, .. }))] =
+        &mut *outer_body.body
+    else {
+        bail!("Expected one inner if statement");
+    };
+
+    outer_if.test = Expression::BooleanOperation(Box::new(BooleanOperation {
+        left: Box::new(parenthesize_and_operand(outer_if.test.clone())),
+        operator: BooleanOp::And {
+            whitespace_before: ParenthesizableWhitespace::SimpleWhitespace(SimpleWhitespace(" ")),
+            whitespace_after: ParenthesizableWhitespace::SimpleWhitespace(SimpleWhitespace(" ")),
+        },
+        right: Box::new(parenthesize_and_operand(inner_if.test.clone())),
+        lpar: vec![],
+        rpar: vec![],
+    }));
+    outer_if.body = inner_if.body.clone();
+
+    let mut state = CodegenState::default();
+    tree.codegen(&mut state);
+
+    // Reconstruct and reformat the code.
+    let module_text = state.to_string();
+    let module_text = if outer_indent.is_empty() {
+        &module_text
+    } else {
+        module_text.strip_prefix("def f():\n").unwrap()
+    };
+    let contents = if is_elif {
+        module_text.replacen("if", "elif", 1)
+    } else {
+        module_text.to_string()
+    };
+
+    Ok(Fix::replacement(
+        contents,
+        Location::new(stmt.location.row(), 0),
+        Location::new(stmt.end_location.unwrap().row() + 1, 0),
+    ))
+}

src/rules/flake8_simplify/rules/mod.rs

@@ -25,6 +25,7 @@ mod ast_if;
 mod ast_ifexp;
 mod ast_unary_op;
 mod ast_with;
+mod fix_if;
 mod key_in_dict;
 mod open_file_with_context_handler;
 mod return_in_try_except_finally;

src/rules/flake8_simplify/snapshots/ruff__rules__flake8_simplify__tests__SIM102_SIM102.py.snap

@@ -5,21 +5,130 @@ expression: diagnostics
 - kind:
     NestedIfStatements: ~
   location:
-    row: 1
+    row: 2
     column: 0
   end_location:
-    row: 3
+    row: 4
+    column: 9
+  fix:
+    content: "if a and b:\n    c\n"
+    location:
+      row: 2
+      column: 0
+    end_location:
+      row: 5
+      column: 0
+  parent: ~
+- kind:
+    NestedIfStatements: ~
+  location:
+    row: 9
+    column: 0
+  end_location:
+    row: 11
+    column: 9
+  fix:
+    content: "elif b and c:\n    d\n"
+    location:
+      row: 9
+      column: 0
+    end_location:
+      row: 12
+      column: 0
+  parent: ~
+- kind:
+    NestedIfStatements: ~
+  location:
+    row: 14
+    column: 0
+  end_location:
+    row: 17
     column: 9
   fix: ~
   parent: ~
 - kind:
     NestedIfStatements: ~
   location:
-    row: 8
+    row: 20
     column: 0
   end_location:
-    row: 10
+    row: 23
     column: 9
-  fix: ~
+  fix:
+    content: "if a and b:\n    # Fixable due to placement of this comment.\n    c\n"
+    location:
+      row: 20
+      column: 0
+    end_location:
+      row: 24
+      column: 0
+  parent: ~
+- kind:
+    NestedIfStatements: ~
+  location:
+    row: 45
+    column: 4
+  end_location:
+    row: 57
+    column: 23
+  fix:
+    content: "    if True and True:\n        \"\"\"this\nis valid\"\"\"\n\n        \"\"\"the indentation on\n            this line is significant\"\"\"\n\n        \"this is\" \\\n\"allowed too\"\n\n        (\"so is\"\n\"this for some reason\")\n"
+    location:
+      row: 45
+      column: 0
+    end_location:
+      row: 58
+      column: 0
+  parent: ~
+- kind:
+    NestedIfStatements: ~
+  location:
+    row: 61
+    column: 0
+  end_location:
+    row: 73
+    column: 23
+  fix:
+    content: "if True and True:\n    \"\"\"this\nis valid\"\"\"\n\n    \"\"\"the indentation on\n        this line is significant\"\"\"\n\n    \"this is\" \\\n\"allowed too\"\n\n    (\"so is\"\n\"this for some reason\")\n"
+    location:
+      row: 61
+      column: 0
+    end_location:
+      row: 74
+      column: 0
+  parent: ~
+- kind:
+    NestedIfStatements: ~
+  location:
+    row: 77
+    column: 4
+  end_location:
+    row: 81
+    column: 32
+  fix:
+    content: "    if node.module and (node.module == \"multiprocessing\" or node.module.startswith(\n        \"multiprocessing.\"\n    )):\n        print(\"Bad module!\")\n"
+    location:
+      row: 77
+      column: 0
+    end_location:
+      row: 82
+      column: 0
+  parent: ~
+- kind:
+    NestedIfStatements: ~
+  location:
+    row: 84
+    column: 0
+  end_location:
+    row: 88
+    column: 28
+  fix:
+    content: "if node.module and (node.module == \"multiprocessing\" or node.module.startswith(\n    \"multiprocessing.\"\n)):\n    print(\"Bad module!\")\n"
+    location:
+      row: 84
+      column: 0
+    end_location:
+      row: 89
+      column: 0
   parent: ~
 

src/violations.rs

@@ -2829,11 +2829,15 @@ impl AlwaysAutofixableViolation for DuplicateIsinstanceCall {
 define_violation!(
     pub struct NestedIfStatements;
 );
-impl Violation for NestedIfStatements {
+impl AlwaysAutofixableViolation for NestedIfStatements {
     fn message(&self) -> String {
         "Use a single `if` statement instead of nested `if` statements".to_string()
     }
 
+    fn autofix_title(&self) -> String {
+        "Combine `if` statements using `and`".to_string()
+    }
+
     fn placeholder() -> Self {
         NestedIfStatements
     }