Respect async with in timeout-without-await (#9859)

Closes https://github.com/astral-sh/ruff/issues/9855.

crates/ruff_linter/resources/test/fixtures/flake8_trio/TRIO100.py

@@ -1,18 +1,27 @@
 import trio
 
 
-async def foo():
+async def func():
     with trio.fail_after():
         ...
 
-async def foo():
+
+async def func():
     with trio.fail_at():
         await ...
 
-async def foo():
+
+async def func():
     with trio.move_on_after():
         ...
 
-async def foo():
+
+async def func():
     with trio.move_at():
         await ...
+
+
+async def func():
+    with trio.move_at():
+        async with trio.open_nursery() as nursery:
+            ...

crates/ruff_linter/src/rules/flake8_trio/snapshots/ruff_linter__rules__flake8_trio__tests__TRIO100_TRIO100.py.snap

@@ -3,24 +3,20 @@ source: crates/ruff_linter/src/rules/flake8_trio/mod.rs
 ---
 TRIO100.py:5:5: TRIO100 A `with trio.fail_after(...):` context does not contain any `await` statements. This makes it pointless, as the timeout can only be triggered by a checkpoint.
   |
-4 |   async def foo():
+4 |   async def func():
 5 |       with trio.fail_after():
   |  _____^
 6 | |         ...
   | |___________^ TRIO100
-7 |   
-8 |   async def foo():
   |
 
-TRIO100.py:13:5: TRIO100 A `with trio.move_on_after(...):` context does not contain any `await` statements. This makes it pointless, as the timeout can only be triggered by a checkpoint.
+TRIO100.py:15:5: TRIO100 A `with trio.move_on_after(...):` context does not contain any `await` statements. This makes it pointless, as the timeout can only be triggered by a checkpoint.
    |
-12 |   async def foo():
-13 |       with trio.move_on_after():
+14 |   async def func():
+15 |       with trio.move_on_after():
    |  _____^
-14 | |         ...
+16 | |         ...
    | |___________^ TRIO100
-15 |   
-16 |   async def foo():
    |
 
 

crates/ruff_python_ast/src/helpers.rs

@@ -1021,6 +1021,12 @@ impl Visitor<'_> for AwaitVisitor {
     fn visit_stmt(&mut self, stmt: &Stmt) {
         match stmt {
             Stmt::FunctionDef(_) | Stmt::ClassDef(_) => (),
+            Stmt::With(ast::StmtWith { is_async: true, .. }) => {
+                self.seen_await = true;
+            }
+            Stmt::For(ast::StmtFor { is_async: true, .. }) => {
+                self.seen_await = true;
+            }
             _ => crate::visitor::walk_stmt(self, stmt),
         }
     }