Autofix SIM117 (MultipleWithStatements) (#1961)

This is slightly buggy due to Instagram/LibCST#855; it will complain `[ERROR] Failed to fix nested with: Failed to extract CST from source` when trying to fix nested parenthesized `with` statements lacking trailing commas. But presumably people who write parenthesized `with` statements already knew that they don’t need to nest them.

Signed-off-by: Anders Kaseorg <andersk@mit.edu>

README.md

@@ -1014,7 +1014,7 @@ For more, see [flake8-simplify](https://pypi.org/project/flake8-simplify/0.19.3/
 | SIM111 | ConvertLoopToAll | Use `return all(x for x in y)` instead of `for` loop | 🛠 |
 | SIM112 | UseCapitalEnvironmentVariables | Use capitalized environment variable `...` instead of `...` | 🛠 |
 | SIM115 | OpenFileWithContextHandler | Use context handler for opening files |  |
-| SIM117 | MultipleWithStatements | Use a single `with` statement with multiple contexts instead of nested `with` statements |  |
+| SIM117 | MultipleWithStatements | Use a single `with` statement with multiple contexts instead of nested `with` statements | 🛠 |
 | SIM118 | KeyInDict | Use `key in dict` instead of `key in dict.keys()` | 🛠 |
 | SIM201 | NegateEqualOp | Use `left != right` instead of `not left == right` | 🛠 |
 | SIM202 | NegateNotEqualOp | Use `left == right` instead of `not left != right` | 🛠 |

resources/test/fixtures/flake8_simplify/SIM117.py

@@ -1,30 +1,92 @@
-with A() as a:  # SIM117
+# SIM117
+with A() as a:
     with B() as b:
         print("hello")
 
-with A():  # SIM117
+# SIM117
+with A():
     with B():
         with C():
             print("hello")
 
+# SIM117
+with A() as a:
+    # Unfixable due to placement of this comment.
+    with B() as b:
+        print("hello")
+
+# SIM117
+with A() as a:
+    with B() as b:
+        # Fixable due to placement of this comment.
+        print("hello")
+
+# OK
 with A() as a:
     a()
     with B() as b:
         print("hello")
 
+# OK
 with A() as a:
     with B() as b:
         print("hello")
     a()
 
+# OK
 async with A() as a:
     with B() as b:
         print("hello")
 
+# OK
 with A() as a:
     async with B() as b:
         print("hello")
 
+# OK
 async with A() as a:
     async with B() as b:
         print("hello")
+
+while True:
+    # SIM117
+    with A() as a:
+        with B() as b:
+            """this
+is valid"""
+
+            """the indentation on
+            this line is significant"""
+
+            "this is" \
+"allowed too"
+
+            ("so is"
+"this for some reason")
+
+# SIM117
+with (
+    A() as a,
+    B() as b,
+):
+    with C() as c:
+        print("hello")
+
+# SIM117
+with A() as a:
+    with (
+        B() as b,
+        C() as c,
+    ):
+        print("hello")
+
+# SIM117
+with (
+    A() as a,
+    B() as b,
+):
+    with (
+        C() as c,
+        D() as d,
+    ):
+        print("hello")

src/rules/flake8_simplify/rules/ast_with.rs

@@ -1,9 +1,11 @@
+use log::error;
 use rustpython_ast::{Located, Stmt, StmtKind, Withitem};
 
-use crate::ast::helpers::first_colon_range;
+use super::fix_with;
+use crate::ast::helpers::{first_colon_range, has_comments_in};
 use crate::ast::types::Range;
 use crate::checkers::ast::Checker;
-use crate::registry::Diagnostic;
+use crate::registry::{Diagnostic, RuleCode};
 use crate::violations;
 
 fn find_last_with(body: &[Stmt]) -> Option<(&Vec<Withitem>, &Vec<Stmt>)> {
@@ -40,12 +42,33 @@ pub fn multiple_with_statements(
             ),
             checker.locator,
         );
-        checker.diagnostics.push(Diagnostic::new(
+        let mut diagnostic = Diagnostic::new(
             violations::MultipleWithStatements,
             colon.map_or_else(
                 || Range::from_located(with_stmt),
                 |colon| Range::new(with_stmt.location, colon.end_location),
             ),
-        ));
+        );
+        if checker.patch(&RuleCode::SIM117) {
+            let nested_with = &with_body[0];
+            if !has_comments_in(
+                Range::new(with_stmt.location, nested_with.location),
+                checker.locator,
+            ) {
+                match fix_with::fix_multiple_with_statements(checker.locator, with_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 with: {err}"),
+                }
+            }
+        }
+        checker.diagnostics.push(diagnostic);
     }
 }

src/rules/flake8_simplify/rules/fix_with.rs

@@ -0,0 +1,96 @@
+use std::borrow::Cow;
+
+use anyhow::{bail, Result};
+use libcst_native::{Codegen, CodegenState, CompoundStatement, Statement, Suite, With};
+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;
+
+/// (SIM117) Convert `with a: with b:` to `with a, b:`.
+pub(crate) fn fix_multiple_with_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),
+    ));
+
+    // 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() {
+        contents
+    } else {
+        Cow::Owned(format!("def f():\n{contents}"))
+    };
+
+    // 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::With(outer_with))] = statements else {
+        bail!("Expected one outer with statement")
+    };
+
+    let With {
+        body: Suite::IndentedBlock(ref mut outer_body),
+        ..
+    } = outer_with else {
+        bail!("Expected outer with to have indented body")
+    };
+
+    let [Statement::Compound(CompoundStatement::With(inner_with))] =
+        &mut *outer_body.body
+    else {
+        bail!("Expected one inner with statement");
+    };
+
+    outer_with.items.append(&mut inner_with.items);
+    if outer_with.lpar.is_none() {
+        outer_with.lpar = inner_with.lpar.clone();
+        outer_with.rpar = inner_with.rpar.clone();
+    }
+    outer_with.body = inner_with.body.clone();
+
+    let mut state = CodegenState::default();
+    tree.codegen(&mut state);
+
+    // Reconstruct and reformat the code.
+    let module_text = state.to_string();
+    let contents = if outer_indent.is_empty() {
+        module_text
+    } else {
+        module_text.strip_prefix("def f():\n").unwrap().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

@@ -26,6 +26,7 @@ mod ast_ifexp;
 mod ast_unary_op;
 mod ast_with;
 mod fix_if;
+mod fix_with;
 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__SIM117_SIM117.py.snap

@@ -5,21 +5,130 @@ expression: diagnostics
 - kind:
     MultipleWithStatements: ~
   location:
-    row: 1
+    row: 2
     column: 0
   end_location:
-    row: 2
+    row: 3
     column: 18
-  fix: ~
+  fix:
+    content: "with A() as a, B() as b:\n    print(\"hello\")\n"
+    location:
+      row: 2
+      column: 0
+    end_location:
+      row: 5
+      column: 0
   parent: ~
 - kind:
     MultipleWithStatements: ~
   location:
-    row: 5
+    row: 7
     column: 0
   end_location:
-    row: 7
+    row: 9
     column: 17
+  fix:
+    content: "with A(), B():\n    with C():\n        print(\"hello\")\n"
+    location:
+      row: 7
+      column: 0
+    end_location:
+      row: 11
+      column: 0
+  parent: ~
+- kind:
+    MultipleWithStatements: ~
+  location:
+    row: 13
+    column: 0
+  end_location:
+    row: 15
+    column: 18
   fix: ~
   parent: ~
+- kind:
+    MultipleWithStatements: ~
+  location:
+    row: 19
+    column: 0
+  end_location:
+    row: 20
+    column: 18
+  fix:
+    content: "with A() as a, B() as b:\n    # Fixable due to placement of this comment.\n    print(\"hello\")\n"
+    location:
+      row: 19
+      column: 0
+    end_location:
+      row: 23
+      column: 0
+  parent: ~
+- kind:
+    MultipleWithStatements: ~
+  location:
+    row: 53
+    column: 4
+  end_location:
+    row: 54
+    column: 22
+  fix:
+    content: "    with A() as a, B() as b:\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: 53
+      column: 0
+    end_location:
+      row: 66
+      column: 0
+  parent: ~
+- kind:
+    MultipleWithStatements: ~
+  location:
+    row: 68
+    column: 0
+  end_location:
+    row: 72
+    column: 18
+  fix:
+    content: "with (\n    A() as a,\n    B() as b,C() as c\n):\n    print(\"hello\")\n"
+    location:
+      row: 68
+      column: 0
+    end_location:
+      row: 74
+      column: 0
+  parent: ~
+- kind:
+    MultipleWithStatements: ~
+  location:
+    row: 76
+    column: 0
+  end_location:
+    row: 80
+    column: 6
+  fix:
+    content: "with (\n    A() as a, B() as b,\n    C() as c,\n):\n    print(\"hello\")\n"
+    location:
+      row: 76
+      column: 0
+    end_location:
+      row: 82
+      column: 0
+  parent: ~
+- kind:
+    MultipleWithStatements: ~
+  location:
+    row: 84
+    column: 0
+  end_location:
+    row: 91
+    column: 6
+  fix:
+    content: "with (\n    A() as a,\n    B() as b,C() as c,\n    D() as d,\n):\n    print(\"hello\")\n"
+    location:
+      row: 84
+      column: 0
+    end_location:
+      row: 93
+      column: 0
+  parent: ~