Fix miscellaneous issues in await-outside-async detection (#14218)

## Summary

Closes #14167.

crates/ruff_linter/resources/test/fixtures/pylint/await_outside_async.py

@@ -1,20 +1,24 @@
 # pylint: disable=missing-docstring,unused-variable
 import asyncio
 
+
 async def nested():
     return 42
 
+
 async def main():
     nested()
     print(await nested())  # This is okay
 
+
 def not_async():
     print(await nested())  # [await-outside-async]
 
 
 async def func(i):
     return i**2
 
+
 async def okay_function():
     var = [await func(i) for i in range(5)]  # This should be okay
 
@@ -28,3 +32,43 @@ def inner_func():
 def outer_func():
     async def inner_func():
         await asyncio.sleep(1)
+
+
+def async_for_loop():
+    async for x in foo():
+        pass
+
+
+def async_with():
+    async with foo():
+        pass
+
+
+# See: https://github.com/astral-sh/ruff/issues/14167
+def async_for_generator_elt():
+    (await x for x in foo())
+
+
+# See: https://github.com/astral-sh/ruff/issues/14167
+def async_for_list_comprehension_elt():
+    [await x for x in foo()]
+
+
+# See: https://github.com/astral-sh/ruff/issues/14167
+def async_for_list_comprehension():
+    [x async for x in foo()]
+
+
+# See: https://github.com/astral-sh/ruff/issues/14167
+def await_generator_iter():
+    (x for x in await foo())
+
+
+# See: https://github.com/astral-sh/ruff/issues/14167
+def await_generator_target():
+    (x async for x in foo())
+
+
+# See: https://github.com/astral-sh/ruff/issues/14167
+def async_for_list_comprehension_target():
+    [x for x in await foo()]

crates/ruff_linter/src/checkers/ast/analyze/comprehension.rs

@@ -2,7 +2,7 @@ use ruff_python_ast::Comprehension;
 
 use crate::checkers::ast::Checker;
 use crate::codes::Rule;
-use crate::rules::{flake8_simplify, refurb};
+use crate::rules::{flake8_simplify, pylint, refurb};
 
 /// Run lint rules over a [`Comprehension`] syntax nodes.
 pub(crate) fn comprehension(comprehension: &Comprehension, checker: &mut Checker) {
@@ -12,4 +12,9 @@ pub(crate) fn comprehension(comprehension: &Comprehension, checker: &mut Checker
     if checker.enabled(Rule::ReadlinesInFor) {
         refurb::rules::readlines_in_comprehension(checker, comprehension);
     }
+    if comprehension.is_async {
+        if checker.enabled(Rule::AwaitOutsideAsync) {
+            pylint::rules::await_outside_async(checker, comprehension);
+        }
+    }
 }

crates/ruff_linter/src/checkers/ast/analyze/statement.rs

@@ -1303,7 +1303,14 @@ pub(crate) fn statement(stmt: &Stmt, checker: &mut Checker) {
                 ruff::rules::assert_with_print_message(checker, assert_stmt);
             }
         }
-        Stmt::With(with_stmt @ ast::StmtWith { items, body, .. }) => {
+        Stmt::With(
+            with_stmt @ ast::StmtWith {
+                items,
+                body,
+                is_async,
+                ..
+            },
+        ) => {
             if checker.enabled(Rule::TooManyNestedBlocks) {
                 pylint::rules::too_many_nested_blocks(checker, stmt);
             }
@@ -1335,6 +1342,11 @@ pub(crate) fn statement(stmt: &Stmt, checker: &mut Checker) {
             if checker.enabled(Rule::CancelScopeNoCheckpoint) {
                 flake8_async::rules::cancel_scope_no_checkpoint(checker, with_stmt, items);
             }
+            if *is_async {
+                if checker.enabled(Rule::AwaitOutsideAsync) {
+                    pylint::rules::await_outside_async(checker, stmt);
+                }
+            }
         }
         Stmt::While(while_stmt @ ast::StmtWhile { body, orelse, .. }) => {
             if checker.enabled(Rule::TooManyNestedBlocks) {
@@ -1422,7 +1434,11 @@ pub(crate) fn statement(stmt: &Stmt, checker: &mut Checker) {
             if checker.enabled(Rule::ReadlinesInFor) {
                 refurb::rules::readlines_in_for(checker, for_stmt);
             }
-            if !is_async {
+            if *is_async {
+                if checker.enabled(Rule::AwaitOutsideAsync) {
+                    pylint::rules::await_outside_async(checker, stmt);
+                }
+            } else {
                 if checker.enabled(Rule::ReimplementedBuiltin) {
                     flake8_simplify::rules::convert_for_loop_to_any_all(checker, stmt);
                 }

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

@@ -53,9 +53,9 @@ use ruff_python_parser::{Parsed, Tokens};
 use ruff_python_semantic::all::{DunderAllDefinition, DunderAllFlags};
 use ruff_python_semantic::analyze::{imports, typing};
 use ruff_python_semantic::{
-    BindingFlags, BindingId, BindingKind, Exceptions, Export, FromImport, Globals, Import, Module,
-    ModuleKind, ModuleSource, NodeId, ScopeId, ScopeKind, SemanticModel, SemanticModelFlags,
-    StarImport, SubmoduleImport,
+    BindingFlags, BindingId, BindingKind, Exceptions, Export, FromImport, GeneratorKind, Globals,
+    Import, Module, ModuleKind, ModuleSource, NodeId, ScopeId, ScopeKind, SemanticModel,
+    SemanticModelFlags, StarImport, SubmoduleImport,
 };
 use ruff_python_stdlib::builtins::{python_builtins, MAGIC_GLOBALS};
 use ruff_python_trivia::CommentRanges;
@@ -1137,19 +1137,25 @@ impl<'a> Visitor<'a> for Checker<'a> {
                 elt,
                 generators,
                 range: _,
-            })
-            | Expr::SetComp(ast::ExprSetComp {
+            }) => {
+                self.visit_generators(GeneratorKind::ListComprehension, generators);
+                self.visit_expr(elt);
+            }
+            Expr::SetComp(ast::ExprSetComp {
                 elt,
                 generators,
                 range: _,
-            })
-            | Expr::Generator(ast::ExprGenerator {
+            }) => {
+                self.visit_generators(GeneratorKind::SetComprehension, generators);
+                self.visit_expr(elt);
+            }
+            Expr::Generator(ast::ExprGenerator {
                 elt,
                 generators,
                 range: _,
                 parenthesized: _,
             }) => {
-                self.visit_generators(generators);
+                self.visit_generators(GeneratorKind::Generator, generators);
                 self.visit_expr(elt);
             }
             Expr::DictComp(ast::ExprDictComp {
@@ -1158,7 +1164,7 @@ impl<'a> Visitor<'a> for Checker<'a> {
                 generators,
                 range: _,
             }) => {
-                self.visit_generators(generators);
+                self.visit_generators(GeneratorKind::DictComprehension, generators);
                 self.visit_expr(key);
                 self.visit_expr(value);
             }
@@ -1748,7 +1754,7 @@ impl<'a> Checker<'a> {
 
     /// Visit a list of [`Comprehension`] nodes, assumed to be the comprehensions that compose a
     /// generator expression, like a list or set comprehension.
-    fn visit_generators(&mut self, generators: &'a [Comprehension]) {
+    fn visit_generators(&mut self, kind: GeneratorKind, generators: &'a [Comprehension]) {
         let mut iterator = generators.iter();
 
         let Some(generator) = iterator.next() else {
@@ -1785,7 +1791,7 @@ impl<'a> Checker<'a> {
         // while all subsequent reads and writes are evaluated in the inner scope. In particular,
         // `x` is local to `foo`, and the `T` in `y=T` skips the class scope when resolving.
         self.visit_expr(&generator.iter);
-        self.semantic.push_scope(ScopeKind::Generator);
+        self.semantic.push_scope(ScopeKind::Generator(kind));
 
         self.semantic.flags = flags | SemanticModelFlags::COMPREHENSION_ASSIGNMENT;
         self.visit_expr(&generator.target);

crates/ruff_linter/src/rules/pylint/rules/await_outside_async.rs

@@ -1,16 +1,15 @@
-use ruff_python_ast::Expr;
-
 use ruff_diagnostics::{Diagnostic, Violation};
 use ruff_macros::{derive_message_formats, violation};
+use ruff_python_semantic::{GeneratorKind, ScopeKind};
 use ruff_text_size::Ranged;
 
 use crate::checkers::ast::Checker;
 
 /// ## What it does
-/// Checks for uses of `await` outside of `async` functions.
+/// Checks for uses of `await` outside `async` functions.
 ///
 /// ## Why is this bad?
-/// Using `await` outside of an `async` function is a syntax error.
+/// Using `await` outside an `async` function is a syntax error.
 ///
 /// ## Example
 /// ```python
@@ -44,10 +43,30 @@ impl Violation for AwaitOutsideAsync {
 }
 
 /// PLE1142
-pub(crate) fn await_outside_async(checker: &mut Checker, expr: &Expr) {
-    if !checker.semantic().in_async_context() {
-        checker
-            .diagnostics
-            .push(Diagnostic::new(AwaitOutsideAsync, expr.range()));
+pub(crate) fn await_outside_async<T: Ranged>(checker: &mut Checker, node: T) {
+    // If we're in an `async` function, we're good.
+    if checker.semantic().in_async_context() {
+        return;
+    }
+
+    // Generators are evaluated lazily, so you can use `await` in them. For example:
+    // ```python
+    // # This is valid
+    // (await x for x in y)
+    // (x async for x in y)
+    //
+    // # This is invalid
+    // (x for x in async y)
+    // [await x for x in y]
+    // ```
+    if matches!(
+        checker.semantic().current_scope().kind,
+        ScopeKind::Generator(GeneratorKind::Generator)
+    ) {
+        return;
     }
+
+    checker
+        .diagnostics
+        .push(Diagnostic::new(AwaitOutsideAsync, node.range()));
 }

crates/ruff_linter/src/rules/pylint/snapshots/ruff_linter__rules__pylint__tests__PLE1142_await_outside_async.py.snap

@@ -1,19 +1,67 @@
 ---
 source: crates/ruff_linter/src/rules/pylint/mod.rs
 ---
-await_outside_async.py:12:11: PLE1142 `await` should be used within an async function
+await_outside_async.py:15:11: PLE1142 `await` should be used within an async function
    |
-11 | def not_async():
-12 |     print(await nested())  # [await-outside-async]
+14 | def not_async():
+15 |     print(await nested())  # [await-outside-async]
    |           ^^^^^^^^^^^^^^ PLE1142
    |
 
-await_outside_async.py:25:9: PLE1142 `await` should be used within an async function
+await_outside_async.py:29:9: PLE1142 `await` should be used within an async function
    |
-23 | async def func2():
-24 |     def inner_func():
-25 |         await asyncio.sleep(1)  # [await-outside-async]
+27 | async def func2():
+28 |     def inner_func():
+29 |         await asyncio.sleep(1)  # [await-outside-async]
    |         ^^^^^^^^^^^^^^^^^^^^^^ PLE1142
    |
 
+await_outside_async.py:38:5: PLE1142 `await` should be used within an async function
+   |
+37 |   def async_for_loop():
+38 |       async for x in foo():
+   |  _____^
+39 | |         pass
+   | |____________^ PLE1142
+   |
+
+await_outside_async.py:43:5: PLE1142 `await` should be used within an async function
+   |
+42 |   def async_with():
+43 |       async with foo():
+   |  _____^
+44 | |         pass
+   | |____________^ PLE1142
+   |
+
+await_outside_async.py:54:6: PLE1142 `await` should be used within an async function
+   |
+52 | # See: https://github.com/astral-sh/ruff/issues/14167
+53 | def async_for_list_comprehension_elt():
+54 |     [await x for x in foo()]
+   |      ^^^^^^^ PLE1142
+   |
 
+await_outside_async.py:59:8: PLE1142 `await` should be used within an async function
+   |
+57 | # See: https://github.com/astral-sh/ruff/issues/14167
+58 | def async_for_list_comprehension():
+59 |     [x async for x in foo()]
+   |        ^^^^^^^^^^^^^^^^^^^^ PLE1142
+   |
+
+await_outside_async.py:64:17: PLE1142 `await` should be used within an async function
+   |
+62 | # See: https://github.com/astral-sh/ruff/issues/14167
+63 | def await_generator_iter():
+64 |     (x for x in await foo())
+   |                 ^^^^^^^^^^^ PLE1142
+   |
+
+await_outside_async.py:74:17: PLE1142 `await` should be used within an async function
+   |
+72 | # See: https://github.com/astral-sh/ruff/issues/14167
+73 | def async_for_list_comprehension_target():
+74 |     [x for x in await foo()]
+   |                 ^^^^^^^^^^^ PLE1142
+   |

crates/ruff_python_semantic/src/scope.rs

@@ -170,13 +170,21 @@ bitflags! {
 pub enum ScopeKind<'a> {
     Class(&'a ast::StmtClassDef),
     Function(&'a ast::StmtFunctionDef),
-    Generator,
+    Generator(GeneratorKind),
     Module,
     /// A Python 3.12+ [annotation scope](https://docs.python.org/3/reference/executionmodel.html#annotation-scopes)
     Type,
     Lambda(&'a ast::ExprLambda),
 }
 
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+pub enum GeneratorKind {
+    Generator,
+    ListComprehension,
+    DictComprehension,
+    SetComprehension,
+}
+
 /// Id uniquely identifying a scope in a program.
 ///
 /// Using a `u32` is sufficient because Ruff only supports parsing documents with a size of max `u32::max`