[syntax-errors] await outside async functions (#17363)

Summary
--

This PR implements detecting the use of `await` expressions outside of
async functions. This is a reimplementation of
[await-outside-async
(PLE1142)](https://docs.astral.sh/ruff/rules/await-outside-async/) as a
semantic syntax error.

Despite the rule name, PLE1142 also applies to `async for` and `async
with`, so these are covered here too.

Test Plan
--

Existing PLE1142 tests.

I also deleted more code from the `SemanticSyntaxCheckerVisitor` to
avoid changes in other parser tests.

crates/ruff_python_parser/src/semantic_errors.rs

@@ -103,12 +103,31 @@ impl SemanticSyntaxChecker {
                     Self::add_error(ctx, SemanticSyntaxErrorKind::ReturnOutsideFunction, *range);
                 }
             }
-            Stmt::For(ast::StmtFor { target, iter, .. }) => {
+            Stmt::For(ast::StmtFor {
+                target,
+                iter,
+                is_async,
+                ..
+            }) => {
                 // test_err single_star_for
                 // for _ in *x: ...
                 // for *x in xs: ...
                 Self::invalid_star_expression(target, ctx);
                 Self::invalid_star_expression(iter, ctx);
+                if *is_async {
+                    Self::await_outside_async_function(
+                        ctx,
+                        stmt,
+                        AwaitOutsideAsyncFunctionKind::AsyncFor,
+                    );
+                }
+            }
+            Stmt::With(ast::StmtWith { is_async: true, .. }) => {
+                Self::await_outside_async_function(
+                    ctx,
+                    stmt,
+                    AwaitOutsideAsyncFunctionKind::AsyncWith,
+                );
             }
             _ => {}
         }
@@ -514,11 +533,13 @@ impl SemanticSyntaxChecker {
             }) => {
                 Self::check_generator_expr(elt, generators, ctx);
                 Self::async_comprehension_outside_async_function(ctx, generators);
-            }
-            Expr::Generator(ast::ExprGenerator {
-                elt, generators, ..
-            }) => {
-                Self::check_generator_expr(elt, generators, ctx);
+                for generator in generators.iter().filter(|g| g.is_async) {
+                    Self::await_outside_async_function(
+                        ctx,
+                        generator,
+                        AwaitOutsideAsyncFunctionKind::AsyncComprehension,
+                    );
+                }
             }
             Expr::DictComp(ast::ExprDictComp {
                 key,
@@ -529,6 +550,20 @@ impl SemanticSyntaxChecker {
                 Self::check_generator_expr(key, generators, ctx);
                 Self::check_generator_expr(value, generators, ctx);
                 Self::async_comprehension_outside_async_function(ctx, generators);
+                for generator in generators.iter().filter(|g| g.is_async) {
+                    Self::await_outside_async_function(
+                        ctx,
+                        generator,
+                        AwaitOutsideAsyncFunctionKind::AsyncComprehension,
+                    );
+                }
+            }
+            Expr::Generator(ast::ExprGenerator {
+                elt, generators, ..
+            }) => {
+                Self::check_generator_expr(elt, generators, ctx);
+                // Note that `await_outside_async_function` is not called here because generators
+                // are evaluated lazily. See the note in the function for more details.
             }
             Expr::Name(ast::ExprName {
                 range,
@@ -603,11 +638,53 @@ impl SemanticSyntaxChecker {
             }
             Expr::Await(_) => {
                 Self::yield_outside_function(ctx, expr, YieldOutsideFunctionKind::Await);
+                Self::await_outside_async_function(ctx, expr, AwaitOutsideAsyncFunctionKind::Await);
             }
             _ => {}
         }
     }
 
+    /// PLE1142
+    fn await_outside_async_function<Ctx: SemanticSyntaxContext, Node: Ranged>(
+        ctx: &Ctx,
+        node: Node,
+        kind: AwaitOutsideAsyncFunctionKind,
+    ) {
+        if ctx.in_async_context() {
+            return;
+        }
+        // `await` is allowed at the top level of a Jupyter notebook.
+        // See: https://ipython.readthedocs.io/en/stable/interactive/autoawait.html.
+        if ctx.in_module_scope() && ctx.in_notebook() {
+            return;
+        }
+        // Generators are evaluated lazily, so you can use `await` in them. For example:
+        //
+        // ```python
+        // # This is valid
+        // def f():
+        //     (await x for x in y)
+        //     (x async for x in y)
+        //
+        // # This is invalid
+        // def f():
+        //     (x for x in await y)
+        //     [await x for x in y]
+        // ```
+        //
+        // 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() {
+            return;
+        }
+        Self::add_error(
+            ctx,
+            SemanticSyntaxErrorKind::AwaitOutsideAsyncFunction(kind),
+            node.range(),
+        );
+    }
+
     /// F704
     fn yield_outside_function<Ctx: SemanticSyntaxContext>(
         ctx: &Ctx,
@@ -803,6 +880,9 @@ impl Display for SemanticSyntaxError {
             SemanticSyntaxErrorKind::ReturnOutsideFunction => {
                 f.write_str("`return` statement outside of a function")
             }
+            SemanticSyntaxErrorKind::AwaitOutsideAsyncFunction(kind) => {
+                write!(f, "`{kind}` outside of an asynchronous function")
+            }
         }
     }
 }
@@ -1101,6 +1181,38 @@ pub enum SemanticSyntaxErrorKind {
 
     /// Represents the use of `return` outside of a function scope.
     ReturnOutsideFunction,
+
+    /// Represents the use of `await`, `async for`, or `async with` outside of an asynchronous
+    /// function.
+    ///
+    /// ## Examples
+    ///
+    /// ```python
+    /// def f():
+    ///     await 1                # error
+    ///     async for x in y: ...  # error
+    ///     async with x: ...      # error
+    /// ```
+    AwaitOutsideAsyncFunction(AwaitOutsideAsyncFunctionKind),
+}
+
+#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
+pub enum AwaitOutsideAsyncFunctionKind {
+    Await,
+    AsyncFor,
+    AsyncWith,
+    AsyncComprehension,
+}
+
+impl Display for AwaitOutsideAsyncFunctionKind {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        f.write_str(match self {
+            AwaitOutsideAsyncFunctionKind::Await => "await",
+            AwaitOutsideAsyncFunctionKind::AsyncFor => "async for",
+            AwaitOutsideAsyncFunctionKind::AsyncWith => "async with",
+            AwaitOutsideAsyncFunctionKind::AsyncComprehension => "asynchronous comprehension",
+        })
+    }
 }
 
 #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
@@ -1527,6 +1639,12 @@ 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.
+    ///
+    /// Note that this refers to an `Expr::Generator` precisely, not to comprehensions more
+    /// generally.
+    fn in_generator_scope(&self) -> bool;
+
     /// Returns `true` if the source file is a Jupyter notebook.
     fn in_notebook(&self) -> bool;
 

crates/ruff_python_parser/tests/fixtures.rs

@@ -462,7 +462,7 @@ impl<'ast> SourceOrderVisitor<'ast> for ValidateAstVisitor<'ast> {
 
 enum Scope {
     Module,
-    Function { is_async: bool },
+    Function,
     Comprehension { is_async: bool },
     Class,
 }
@@ -529,12 +529,7 @@ impl SemanticSyntaxContext for SemanticSyntaxCheckerVisitor<'_> {
     }
 
     fn in_async_context(&self) -> bool {
-        for scope in &self.scopes {
-            if let Scope::Function { is_async } = scope {
-                return *is_async;
-            }
-        }
-        false
+        true
     }
 
     fn in_sync_comprehension(&self) -> bool {
@@ -561,6 +556,10 @@ impl SemanticSyntaxContext for SemanticSyntaxCheckerVisitor<'_> {
     fn in_await_allowed_context(&self) -> bool {
         true
     }
+
+    fn in_generator_scope(&self) -> bool {
+        true
+    }
 }
 
 impl Visitor<'_> for SemanticSyntaxCheckerVisitor<'_> {
@@ -587,10 +586,8 @@ impl Visitor<'_> for SemanticSyntaxCheckerVisitor<'_> {
                 self.visit_body(body);
                 self.scopes.pop().unwrap();
             }
-            ast::Stmt::FunctionDef(ast::StmtFunctionDef { is_async, .. }) => {
-                self.scopes.push(Scope::Function {
-                    is_async: *is_async,
-                });
+            ast::Stmt::FunctionDef(ast::StmtFunctionDef { .. }) => {
+                self.scopes.push(Scope::Function);
                 ast::visitor::walk_stmt(self, stmt);
                 self.scopes.pop().unwrap();
             }
@@ -604,7 +601,7 @@ impl Visitor<'_> for SemanticSyntaxCheckerVisitor<'_> {
         self.with_semantic_checker(|semantic, context| semantic.visit_expr(expr, context));
         match expr {
             ast::Expr::Lambda(_) => {
-                self.scopes.push(Scope::Function { is_async: false });
+                self.scopes.push(Scope::Function);
                 ast::visitor::walk_expr(self, expr);
                 self.scopes.pop().unwrap();
             }