Suggest combining async with statements (#5022)

## Summary

Previously the rule for SIM117 explicitly ignored `async with`
statements as it would incorrectly suggestion to merge `async with` and
regular `with` statements as reported in issue #1902.

This partially reverts the fix for that (commit
396be5edea) by enabling the rules for
`async with` statements again, but with a check ensuring that the
statements are both of the same kind, i.e. both `async with` or both
(just) `with` statements.

Closes #3025

## Test Plan

Updated and existing test and added a new test case from #3025.

crates/ruff/resources/test/fixtures/flake8_simplify/SIM117.py

@@ -33,17 +33,17 @@ with A() as a:
         print("hello")
     a()
 
-# OK
+# OK, can't merge async with and with.
 async with A() as a:
     with B() as b:
         print("hello")
 
-# OK
+# OK, can't merge async with and with.
 with A() as a:
     async with B() as b:
         print("hello")
 
-# OK
+# SIM117
 async with A() as a:
     async with B() as b:
         print("hello")
@@ -99,4 +99,25 @@ with A("01ß9💣2ℝ8901ß9💣2ℝ8901ß9💣2ℝ89") as a:
 # SIM117 (not auto-fixable too long)
 with A("01ß9💣2ℝ8901ß9💣2ℝ8901ß9💣2ℝ890") as a:
     with B("01ß9💣2ℝ8901ß9💣2ℝ8901ß9💣2ℝ89") as b:
-        print("hello")
+        print("hello")
+
+# From issue #3025.
+async def main():
+    async with A() as a: # SIM117.
+        async with B() as b:
+            print("async-inside!")
+
+    return 0
+
+# OK. Can't merge across different kinds of with statements.
+with a as a2:
+    async with b as b2:
+        with c as c2:
+            async with d as d2:
+                f(a2, b2, c2, d2)
+
+# OK. Can't merge across different kinds of with statements.
+async with b as b2:
+    with c as c2:
+        async with d as d2:
+            f(b2, c2, d2)

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

@@ -1514,7 +1514,8 @@ where
                     pygrep_hooks::rules::non_existent_mock_method(self, test);
                 }
             }
-            Stmt::With(ast::StmtWith { items, body, .. }) => {
+            Stmt::With(ast::StmtWith { items, body, .. })
+            | Stmt::AsyncWith(ast::StmtAsyncWith { items, body, .. }) => {
                 if self.enabled(Rule::AssertRaisesException) {
                     flake8_bugbear::rules::assert_raises_exception(self, stmt, items);
                 }

crates/ruff/src/rules/flake8_simplify/rules/ast_with.rs

@@ -60,9 +60,14 @@ impl Violation for MultipleWithStatements {
     }
 }
 
-fn find_last_with(body: &[Stmt]) -> Option<(&[Withitem], &[Stmt])> {
-    let [Stmt::With(ast::StmtWith { items, body, .. })] = body else { return None };
-    find_last_with(body).or(Some((items, body)))
+/// Returns a boolean indicating whether it's an async with statement, the items
+/// and body.
+fn next_with(body: &[Stmt]) -> Option<(bool, &[Withitem], &[Stmt])> {
+    match body {
+        [Stmt::With(ast::StmtWith { items, body, .. })] => Some((false, items, body)),
+        [Stmt::AsyncWith(ast::StmtAsyncWith { items, body, .. })] => Some((true, items, body)),
+        _ => None,
+    }
 }
 
 /// SIM117
@@ -72,12 +77,38 @@ pub(crate) fn multiple_with_statements(
     with_body: &[Stmt],
     with_parent: Option<&Stmt>,
 ) {
+    // Make sure we fix from top to bottom for nested with statements, e.g. for
+    // ```python
+    // with A():
+    //     with B():
+    //         with C():
+    //             print("hello")
+    // ```
+    // suggests
+    // ```python
+    // with A(), B():
+    //     with C():
+    //         print("hello")
+    // ```
+    // but not the following
+    // ```python
+    // with A():
+    //     with B(), C():
+    //         print("hello")
+    // ```
     if let Some(Stmt::With(ast::StmtWith { body, .. })) = with_parent {
         if body.len() == 1 {
             return;
         }
     }
-    if let Some((items, body)) = find_last_with(with_body) {
+
+    if let Some((is_async, items, body)) = next_with(with_body) {
+        if is_async != with_stmt.is_async_with_stmt() {
+            // One of the statements is an async with, while the other is not,
+            // we can't merge those statements.
+            return;
+        }
+
         let last_item = items.last().expect("Expected items to be non-empty");
         let colon = first_colon_range(
             TextRange::new(

crates/ruff/src/rules/flake8_simplify/snapshots/ruff__rules__flake8_simplify__tests__SIM117_SIM117.py.snap

@@ -27,8 +27,8 @@ SIM117.py:7:1: SIM117 [*] Use a single `with` statement with multiple contexts i
  6 |   # SIM117
  7 | / with A():
  8 | |     with B():
- 9 | |         with C():
-   | |_________________^ SIM117
+   | |_____________^ SIM117
+ 9 |           with C():
 10 |               print("hello")
    |
    = help: Combine `with` statements
@@ -85,6 +85,29 @@ SIM117.py:19:1: SIM117 [*] Use a single `with` statement with multiple contexts
 24 23 | # OK
 25 24 | with A() as a:
 
+SIM117.py:47:1: SIM117 [*] Use a single `with` statement with multiple contexts instead of nested `with` statements
+   |
+46 |   # SIM117
+47 | / async with A() as a:
+48 | |     async with B() as b:
+   | |________________________^ SIM117
+49 |           print("hello")
+   |
+   = help: Combine `with` statements
+
+ℹ Suggested fix
+44 44 |         print("hello")
+45 45 | 
+46 46 | # SIM117
+47    |-async with A() as a:
+48    |-    async with B() as b:
+49    |-        print("hello")
+   47 |+async with A() as a, B() as b:
+   48 |+    print("hello")
+50 49 | 
+51 50 | while True:
+52 51 |     # SIM117
+
 SIM117.py:53:5: SIM117 [*] Use a single `with` statement with multiple contexts instead of nested `with` statements
    |
 51 |   while True:
@@ -249,4 +272,16 @@ SIM117.py:100:1: SIM117 Use a single `with` statement with multiple contexts ins
     |
     = help: Combine `with` statements
 
+SIM117.py:106:5: SIM117 Use a single `with` statement with multiple contexts instead of nested `with` statements
+    |
+104 |   # From issue #3025.
+105 |   async def main():
+106 |       async with A() as a: # SIM117.
+    |  _____^
+107 | |         async with B() as b:
+    | |____________________________^ SIM117
+108 |               print("async-inside!")
+    |
+    = help: Combine `with` statements
+