diff --git a/crates/ruff_linter/resources/test/fixtures/flake8_trio/TRIO100.py b/crates/ruff_linter/resources/test/fixtures/flake8_trio/TRIO100.py index 86beeb95cd7c3..4499657cc2698 100644 --- a/crates/ruff_linter/resources/test/fixtures/flake8_trio/TRIO100.py +++ b/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: + ... diff --git a/crates/ruff_linter/src/rules/flake8_trio/snapshots/ruff_linter__rules__flake8_trio__tests__TRIO100_TRIO100.py.snap b/crates/ruff_linter/src/rules/flake8_trio/snapshots/ruff_linter__rules__flake8_trio__tests__TRIO100_TRIO100.py.snap index ad7e5b0475ae1..1d364b2670695 100644 --- a/crates/ruff_linter/src/rules/flake8_trio/snapshots/ruff_linter__rules__flake8_trio__tests__TRIO100_TRIO100.py.snap +++ b/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(): | diff --git a/crates/ruff_python_ast/src/helpers.rs b/crates/ruff_python_ast/src/helpers.rs index 98370dfa9b1bc..45e480372ee2d 100644 --- a/crates/ruff_python_ast/src/helpers.rs +++ b/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), } }