Fix syntax error false positives for `await` outside functions (#21763)

## Summary

Fixes #21750 and a related bug in `PLE1142`. We were not properly
considering generators to be valid `await` contexts, which caused the
`F704` issue. One of the tests I added for this also uncovered an issue
in `PLE1142` for comprehensions nested within async generators because
we were only checking the current scope rather than traversing the
nested context.

## Test Plan

Both of these rules are implemented as semantic syntax errors, so I
added tests (and fixes) in both Ruff and ty.

crates/ruff_linter/resources/test/fixtures/pyflakes/F704.py

@@ -17,3 +17,24 @@ def _():
 
     # Valid yield scope
     yield 3
+
+
+# await is valid in any generator, sync or async
+(await cor async for cor in f())  # ok
+(await cor for cor in f())  # ok
+
+# but not in comprehensions
+[await cor async for cor in f()]  # F704
+{await cor async for cor in f()}  # F704
+{await cor: 1 async for cor in f()}  # F704
+[await cor for cor in f()]  # F704
+{await cor for cor in f()}  # F704
+{await cor: 1 for cor in f()}  # F704
+
+# or in the iterator of an async generator, which is evaluated in the parent
+# scope
+(cor async for cor in await f())  # F704
+(await cor async for cor in [await c for c in f()])  # F704
+
+# this is also okay because the comprehension is within the generator scope
+([await c for c in cor] async for cor in f())  # ok

crates/ruff_linter/resources/test/fixtures/syntax_errors/await_outside_async_function.py

@@ -3,3 +3,5 @@ def func():
 
 # Top-level await
 await 1
+
+([await c for c in cor] async for cor in func())  # ok

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

@@ -780,6 +780,10 @@ impl SemanticSyntaxContext for Checker<'_> {
             match scope.kind {
                 ScopeKind::Class(_) => return false,
                 ScopeKind::Function(_) | ScopeKind::Lambda(_) => return true,
+                ScopeKind::Generator {
+                    kind: GeneratorKind::Generator,
+                    ..
+                } => return true,
                 ScopeKind::Generator { .. }
                 | ScopeKind::Module
                 | ScopeKind::Type
@@ -829,14 +833,19 @@ impl SemanticSyntaxContext for Checker<'_> {
         self.source_type.is_ipynb()
     }
 
-    fn in_generator_scope(&self) -> bool {
-        matches!(
-            &self.semantic.current_scope().kind,
-            ScopeKind::Generator {
-                kind: GeneratorKind::Generator,
-                ..
+    fn in_generator_context(&self) -> bool {
+        for scope in self.semantic.current_scopes() {
+            if matches!(
+                scope.kind,
+                ScopeKind::Generator {
+                    kind: GeneratorKind::Generator,
+                    ..
+                }
+            ) {
+                return true;
             }
-        )
+        }
+        false
     }
 
     fn in_loop_context(&self) -> bool {

crates/ruff_linter/src/rules/pyflakes/snapshots/ruff_linter__rules__pyflakes__tests__F704_F704.py.snap

@@ -37,3 +37,88 @@ F704 `await` statement outside of a function
 12 |
 13 | def _():
    |
+
+F704 `await` statement outside of a function
+  --> F704.py:27:2
+   |
+26 | # but not in comprehensions
+27 | [await cor async for cor in f()]  # F704
+   |  ^^^^^^^^^
+28 | {await cor async for cor in f()}  # F704
+29 | {await cor: 1 async for cor in f()}  # F704
+   |
+
+F704 `await` statement outside of a function
+  --> F704.py:28:2
+   |
+26 | # but not in comprehensions
+27 | [await cor async for cor in f()]  # F704
+28 | {await cor async for cor in f()}  # F704
+   |  ^^^^^^^^^
+29 | {await cor: 1 async for cor in f()}  # F704
+30 | [await cor for cor in f()]  # F704
+   |
+
+F704 `await` statement outside of a function
+  --> F704.py:29:2
+   |
+27 | [await cor async for cor in f()]  # F704
+28 | {await cor async for cor in f()}  # F704
+29 | {await cor: 1 async for cor in f()}  # F704
+   |  ^^^^^^^^^
+30 | [await cor for cor in f()]  # F704
+31 | {await cor for cor in f()}  # F704
+   |
+
+F704 `await` statement outside of a function
+  --> F704.py:30:2
+   |
+28 | {await cor async for cor in f()}  # F704
+29 | {await cor: 1 async for cor in f()}  # F704
+30 | [await cor for cor in f()]  # F704
+   |  ^^^^^^^^^
+31 | {await cor for cor in f()}  # F704
+32 | {await cor: 1 for cor in f()}  # F704
+   |
+
+F704 `await` statement outside of a function
+  --> F704.py:31:2
+   |
+29 | {await cor: 1 async for cor in f()}  # F704
+30 | [await cor for cor in f()]  # F704
+31 | {await cor for cor in f()}  # F704
+   |  ^^^^^^^^^
+32 | {await cor: 1 for cor in f()}  # F704
+   |
+
+F704 `await` statement outside of a function
+  --> F704.py:32:2
+   |
+30 | [await cor for cor in f()]  # F704
+31 | {await cor for cor in f()}  # F704
+32 | {await cor: 1 for cor in f()}  # F704
+   |  ^^^^^^^^^
+33 |
+34 | # or in the iterator of an async generator, which is evaluated in the parent
+   |
+
+F704 `await` statement outside of a function
+  --> F704.py:36:23
+   |
+34 | # or in the iterator of an async generator, which is evaluated in the parent
+35 | # scope
+36 | (cor async for cor in await f())  # F704
+   |                       ^^^^^^^^^
+37 | (await cor async for cor in [await c for c in f()])  # F704
+   |
+
+F704 `await` statement outside of a function
+  --> F704.py:37:30
+   |
+35 | # scope
+36 | (cor async for cor in await f())  # F704
+37 | (await cor async for cor in [await c for c in f()])  # F704
+   |                              ^^^^^^^
+38 |
+39 | # this is also okay because the comprehension is within the generator scope
+   |

crates/ruff_linter/src/snapshots/ruff_linter__linter__tests__await_outside_async_function.py.snap

@@ -17,4 +17,6 @@ PLE1142 `await` should be used within an async function
 4 | # Top-level await
 5 | await 1
   | ^^^^^^^
+6 |
+7 | ([await c for c in cor] async for cor in func())  # ok
   |

crates/ruff_python_parser/src/semantic_errors.rs

@@ -896,7 +896,7 @@ impl SemanticSyntaxChecker {
         // This check is required in addition to avoiding calling this function in `visit_expr`
         // because the generator scope applies to nested parts of the `Expr::Generator` that are
         // visited separately.
-        if ctx.in_generator_scope() {
+        if ctx.in_generator_context() {
             return;
         }
         Self::add_error(
@@ -2096,11 +2096,11 @@ pub trait SemanticSyntaxContext {
     /// Returns `true` if the visitor is in a function scope.
     fn in_function_scope(&self) -> bool;
 
-    /// Returns `true` if the visitor is in a generator scope.
+    /// Returns `true` if the visitor is within a generator scope.
     ///
     /// Note that this refers to an `Expr::Generator` precisely, not to comprehensions more
     /// generally.
-    fn in_generator_scope(&self) -> bool;
+    fn in_generator_context(&self) -> bool;
 
     /// Returns `true` if the source file is a Jupyter notebook.
     fn in_notebook(&self) -> bool;

crates/ruff_python_parser/tests/fixtures.rs

@@ -573,7 +573,7 @@ impl SemanticSyntaxContext for SemanticSyntaxCheckerVisitor<'_> {
         true
     }
 
-    fn in_generator_scope(&self) -> bool {
+    fn in_generator_context(&self) -> bool {
         true
     }
 

crates/ty_python_semantic/resources/mdtest/diagnostics/semantic_syntax_errors.md

@@ -143,6 +143,10 @@ await C()
 def f():
     # error: [invalid-syntax] "`await` outside of an asynchronous function"
     await C()
+
+(await cor async for cor in f())  # ok
+(await cor for cor in f())  # ok
+([await c for c in cor] async for cor in f())  # ok
 ```
 
 Generators are evaluated lazily, so `await` is allowed, even outside of a function.

crates/ty_python_semantic/src/semantic_index/builder.rs

@@ -2845,6 +2845,11 @@ impl SemanticSyntaxContext for SemanticIndexBuilder<'_, '_> {
             match scope.kind() {
                 ScopeKind::Class => return false,
                 ScopeKind::Function | ScopeKind::Lambda => return true,
+                ScopeKind::Comprehension
+                    if matches!(scope.node(), NodeWithScopeKind::GeneratorExpression(_)) =>
+                {
+                    return true;
+                }
                 ScopeKind::Comprehension
                 | ScopeKind::Module
                 | ScopeKind::TypeAlias
@@ -2894,11 +2899,14 @@ impl SemanticSyntaxContext for SemanticIndexBuilder<'_, '_> {
         matches!(kind, ScopeKind::Function | ScopeKind::Lambda)
     }
 
-    fn in_generator_scope(&self) -> bool {
-        matches!(
-            self.scopes[self.current_scope()].node(),
-            NodeWithScopeKind::GeneratorExpression(_)
-        )
+    fn in_generator_context(&self) -> bool {
+        for scope_info in &self.scope_stack {
+            let scope = &self.scopes[scope_info.file_scope_id];
+            if matches!(scope.node(), NodeWithScopeKind::GeneratorExpression(_)) {
+                return true;
+            }
+        }
+        false
     }
 
     fn in_notebook(&self) -> bool {