From 31ca1c306448a9574830a827f87d09dbf81349f8 Mon Sep 17 00:00:00 2001
From: Sid <27780930+autinerd@users.noreply.github.com>
Date: Mon, 7 Oct 2024 14:25:54 +0200
Subject: [PATCH] [`flake8-async`] allow async generators (`ASYNC100`) (#13639)

<!--
Thank you for contributing to Ruff! To help us out with reviewing,
please consider the following:

- Does this pull request include a summary of the change? (See below.)
- Does this pull request include a descriptive title?
- Does this pull request include references to any relevant issues?
-->

## Summary

Treat async generators as "await" in ASYNC100.

Fixes #13637

## Test Plan

Updated snapshot
---
 .../test/fixtures/flake8_async/ASYNC100.py    |  5 ++
 ...e8_async__tests__ASYNC100_ASYNC100.py.snap | 54 +++++++++----------
 crates/ruff_python_ast/src/helpers.rs         |  8 +++
 3 files changed, 40 insertions(+), 27 deletions(-)

diff --git a/crates/ruff_linter/resources/test/fixtures/flake8_async/ASYNC100.py b/crates/ruff_linter/resources/test/fixtures/flake8_async/ASYNC100.py
index 8434073d22..353bd2b73a 100644
--- a/crates/ruff_linter/resources/test/fixtures/flake8_async/ASYNC100.py
+++ b/crates/ruff_linter/resources/test/fixtures/flake8_async/ASYNC100.py
@@ -36,6 +36,11 @@ async def func():
             ...
 
 
+async def main():
+    async with asyncio.timeout(7):
+        print({i async for i in long_running_range()})
+
+
 async def func():
     with anyio.move_on_after(delay=0.2):
         ...
diff --git a/crates/ruff_linter/src/rules/flake8_async/snapshots/ruff_linter__rules__flake8_async__tests__ASYNC100_ASYNC100.py.snap b/crates/ruff_linter/src/rules/flake8_async/snapshots/ruff_linter__rules__flake8_async__tests__ASYNC100_ASYNC100.py.snap
index 0eca205a5b..5d92713d30 100644
--- a/crates/ruff_linter/src/rules/flake8_async/snapshots/ruff_linter__rules__flake8_async__tests__ASYNC100_ASYNC100.py.snap
+++ b/crates/ruff_linter/src/rules/flake8_async/snapshots/ruff_linter__rules__flake8_async__tests__ASYNC100_ASYNC100.py.snap
@@ -19,28 +19,19 @@ ASYNC100.py:18:5: ASYNC100 A `with trio.move_on_after(...):` context does not co
    | |___________^ ASYNC100
    |
 
-ASYNC100.py:40:5: ASYNC100 A `with anyio.move_on_after(...):` context does not contain any `await` statements. This makes it pointless, as the timeout can only be triggered by a checkpoint.
-   |
-39 |   async def func():
-40 |       with anyio.move_on_after(delay=0.2):
-   |  _____^
-41 | |         ...
-   | |___________^ ASYNC100
-   |
-
-ASYNC100.py:45:5: ASYNC100 A `with anyio.fail_after(...):` context does not contain any `await` statements. This makes it pointless, as the timeout can only be triggered by a checkpoint.
+ASYNC100.py:45:5: ASYNC100 A `with anyio.move_on_after(...):` context does not contain any `await` statements. This makes it pointless, as the timeout can only be triggered by a checkpoint.
    |
 44 |   async def func():
-45 |       with anyio.fail_after():
+45 |       with anyio.move_on_after(delay=0.2):
    |  _____^
 46 | |         ...
    | |___________^ ASYNC100
    |
 
-ASYNC100.py:50:5: ASYNC100 A `with anyio.CancelScope(...):` context does not contain any `await` statements. This makes it pointless, as the timeout can only be triggered by a checkpoint.
+ASYNC100.py:50:5: ASYNC100 A `with anyio.fail_after(...):` context does not contain any `await` statements. This makes it pointless, as the timeout can only be triggered by a checkpoint.
    |
 49 |   async def func():
-50 |       with anyio.CancelScope():
+50 |       with anyio.fail_after():
    |  _____^
 51 | |         ...
    | |___________^ ASYNC100
@@ -49,7 +40,7 @@ ASYNC100.py:50:5: ASYNC100 A `with anyio.CancelScope(...):` context does not con
 ASYNC100.py:55:5: ASYNC100 A `with anyio.CancelScope(...):` context does not contain any `await` statements. This makes it pointless, as the timeout can only be triggered by a checkpoint.
    |
 54 |   async def func():
-55 |       with anyio.CancelScope(), nullcontext():
+55 |       with anyio.CancelScope():
    |  _____^
 56 | |         ...
    | |___________^ ASYNC100
@@ -58,44 +49,53 @@ ASYNC100.py:55:5: ASYNC100 A `with anyio.CancelScope(...):` context does not con
 ASYNC100.py:60:5: ASYNC100 A `with anyio.CancelScope(...):` context does not contain any `await` statements. This makes it pointless, as the timeout can only be triggered by a checkpoint.
    |
 59 |   async def func():
-60 |       with nullcontext(), anyio.CancelScope():
+60 |       with anyio.CancelScope(), nullcontext():
    |  _____^
 61 | |         ...
    | |___________^ ASYNC100
    |
 
-ASYNC100.py:65:5: ASYNC100 A `with asyncio.timeout(...):` context does not contain any `await` statements. This makes it pointless, as the timeout can only be triggered by a checkpoint.
+ASYNC100.py:65:5: ASYNC100 A `with anyio.CancelScope(...):` context does not contain any `await` statements. This makes it pointless, as the timeout can only be triggered by a checkpoint.
    |
 64 |   async def func():
-65 |       async with asyncio.timeout(delay=0.2):
+65 |       with nullcontext(), anyio.CancelScope():
    |  _____^
 66 | |         ...
    | |___________^ ASYNC100
    |
 
-ASYNC100.py:70:5: ASYNC100 A `with asyncio.timeout_at(...):` context does not contain any `await` statements. This makes it pointless, as the timeout can only be triggered by a checkpoint.
+ASYNC100.py:70:5: ASYNC100 A `with asyncio.timeout(...):` context does not contain any `await` statements. This makes it pointless, as the timeout can only be triggered by a checkpoint.
    |
 69 |   async def func():
-70 |       async with asyncio.timeout_at(when=0.2):
+70 |       async with asyncio.timeout(delay=0.2):
    |  _____^
 71 | |         ...
    | |___________^ ASYNC100
    |
 
-ASYNC100.py:80:5: ASYNC100 A `with asyncio.timeout(...):` context does not contain any `await` statements. This makes it pointless, as the timeout can only be triggered by a checkpoint.
+ASYNC100.py:75:5: ASYNC100 A `with asyncio.timeout_at(...):` context does not contain any `await` statements. This makes it pointless, as the timeout can only be triggered by a checkpoint.
    |
-79 |   async def func():
-80 |       async with asyncio.timeout(delay=0.2), asyncio.TaskGroup(), asyncio.timeout(delay=0.2):
+74 |   async def func():
+75 |       async with asyncio.timeout_at(when=0.2):
    |  _____^
-81 | |         ...
+76 | |         ...
    | |___________^ ASYNC100
    |
 
-ASYNC100.py:90:5: ASYNC100 A `with asyncio.timeout(...):` context does not contain any `await` statements. This makes it pointless, as the timeout can only be triggered by a checkpoint.
+ASYNC100.py:85:5: ASYNC100 A `with asyncio.timeout(...):` context does not contain any `await` statements. This makes it pointless, as the timeout can only be triggered by a checkpoint.
    |
-89 |   async def func():
-90 |       async with asyncio.timeout(delay=0.2), asyncio.timeout(delay=0.2):
+84 |   async def func():
+85 |       async with asyncio.timeout(delay=0.2), asyncio.TaskGroup(), asyncio.timeout(delay=0.2):
    |  _____^
-91 | |         ...
+86 | |         ...
+   | |___________^ ASYNC100
+   |
+
+ASYNC100.py:95:5: ASYNC100 A `with asyncio.timeout(...):` context does not contain any `await` statements. This makes it pointless, as the timeout can only be triggered by a checkpoint.
+   |
+94 |   async def func():
+95 |       async with asyncio.timeout(delay=0.2), asyncio.timeout(delay=0.2):
+   |  _____^
+96 | |         ...
    | |___________^ ASYNC100
    |
diff --git a/crates/ruff_python_ast/src/helpers.rs b/crates/ruff_python_ast/src/helpers.rs
index 44b48c2b18..565c67e585 100644
--- a/crates/ruff_python_ast/src/helpers.rs
+++ b/crates/ruff_python_ast/src/helpers.rs
@@ -1004,6 +1004,14 @@ impl Visitor<'_> for AwaitVisitor {
             crate::visitor::walk_expr(self, expr);
         }
     }
+
+    fn visit_comprehension(&mut self, comprehension: &'_ crate::Comprehension) {
+        if comprehension.is_async {
+            self.seen_await = true;
+        } else {
+            crate::visitor::walk_comprehension(self, comprehension);
+        }
+    }
 }
 
 /// Return `true` if a `Stmt` is a docstring.