[syntax-errors] yield, yield from, and await outside functions (#17298)

Summary
--

This PR reimplements [yield-outside-function
(F704)](https://docs.astral.sh/ruff/rules/yield-outside-function/) as a
semantic syntax error. Despite the name, this rule covers `yield from`
and `await` in addition to `yield`.

Test Plan
--

New linter tests, along with the existing F704 test.

---------

Co-authored-by: Dhruv Manilawala <dhruvmanila@gmail.com>

crates/ruff_linter/resources/test/fixtures/syntax_errors/await_scope.ipynb

@@ -0,0 +1,57 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "071bef14-c2b5-48d7-8e01-39f16f06328f",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "await 1  # runtime TypeError but not a syntax error"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "b4eaebd3-0fae-4e4c-9509-19ae6f743a08",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "def _():\n",
+    "    await 1  # SyntaxError: await outside async function"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "bf10b4ac-41c2-4daa-8862-9ee3386667ee",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "class _:\n",
+    "    await 1  # SyntaxError: await outside function"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3 (ipykernel)",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.13.2"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}

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

@@ -0,0 +1,41 @@
+yield  # error
+yield 1  # error
+yield from 1  # error
+await 1  # error
+[(yield x) for x in range(3)]  # error
+
+
+def f():
+    yield  # okay
+    yield 1  # okay
+    yield from 1  # okay
+    await 1  # okay
+
+
+lambda: (yield)  # okay
+lambda: (yield 1)  # okay
+lambda: (yield from 1)  # okay
+lambda: (await 1)  # okay
+
+
+def outer():
+    class C:
+        yield 1  # error
+
+    [(yield 1) for x in range(3)]  # error
+    ((yield 1) for x in range(3))  # error
+    {(yield 1) for x in range(3)}  # error
+    {(yield 1): 0 for x in range(3)}  # error
+    {0: (yield 1) for x in range(3)}  # error
+
+
+async def outer():
+    [await x for x in range(3)]  # okay, comprehensions don't break async scope
+
+    class C:
+        [await x for x in range(3)]  # error, classes break async scope
+
+    lambda x: await x  # okay for now, lambda breaks _async_ scope but is a function
+
+
+await 1  # error

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

@@ -1203,17 +1203,11 @@ pub(crate) fn expression(expr: &Expr, checker: &Checker) {
             }
         }
         Expr::Yield(_) => {
-            if checker.enabled(Rule::YieldOutsideFunction) {
-                pyflakes::rules::yield_outside_function(checker, expr);
-            }
             if checker.enabled(Rule::YieldInInit) {
                 pylint::rules::yield_in_init(checker, expr);
             }
         }
         Expr::YieldFrom(yield_from) => {
-            if checker.enabled(Rule::YieldOutsideFunction) {
-                pyflakes::rules::yield_outside_function(checker, expr);
-            }
             if checker.enabled(Rule::YieldInInit) {
                 pylint::rules::yield_in_init(checker, expr);
             }
@@ -1222,9 +1216,6 @@ pub(crate) fn expression(expr: &Expr, checker: &Checker) {
             }
         }
         Expr::Await(_) => {
-            if checker.enabled(Rule::YieldOutsideFunction) {
-                pyflakes::rules::yield_outside_function(checker, expr);
-            }
             if checker.enabled(Rule::AwaitOutsideAsync) {
                 pylint::rules::await_outside_async(checker, expr);
             }

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

@@ -66,7 +66,7 @@ use crate::importer::{ImportRequest, Importer, ResolutionError};
 use crate::noqa::NoqaMapping;
 use crate::package::PackageRoot;
 use crate::registry::Rule;
-use crate::rules::pyflakes::rules::LateFutureImport;
+use crate::rules::pyflakes::rules::{LateFutureImport, YieldOutsideFunction};
 use crate::rules::pylint::rules::LoadBeforeGlobalDeclaration;
 use crate::rules::{flake8_pyi, flake8_type_checking, pyflakes, pyupgrade};
 use crate::settings::{flags, LinterSettings};
@@ -589,6 +589,14 @@ impl SemanticSyntaxContext for Checker<'_> {
                     ));
                 }
             }
+            SemanticSyntaxErrorKind::YieldOutsideFunction(kind) => {
+                if self.settings.rules.enabled(Rule::YieldOutsideFunction) {
+                    self.report_diagnostic(Diagnostic::new(
+                        YieldOutsideFunction::new(kind),
+                        error.range,
+                    ));
+                }
+            }
             SemanticSyntaxErrorKind::ReboundComprehensionVariable
             | SemanticSyntaxErrorKind::DuplicateTypeParameter
             | SemanticSyntaxErrorKind::MultipleCaseAssignment(_)
@@ -616,7 +624,25 @@ impl SemanticSyntaxContext for Checker<'_> {
     }
 
     fn in_async_context(&self) -> bool {
-        self.semantic.in_async_context()
+        for scope in self.semantic.current_scopes() {
+            match scope.kind {
+                ScopeKind::Class(_) | ScopeKind::Lambda(_) => return false,
+                ScopeKind::Function(ast::StmtFunctionDef { is_async, .. }) => return *is_async,
+                ScopeKind::Generator { .. } | ScopeKind::Module | ScopeKind::Type => {}
+            }
+        }
+        false
+    }
+
+    fn in_await_allowed_context(&self) -> bool {
+        for scope in self.semantic.current_scopes() {
+            match scope.kind {
+                ScopeKind::Class(_) => return false,
+                ScopeKind::Function(_) | ScopeKind::Lambda(_) => return true,
+                ScopeKind::Generator { .. } | ScopeKind::Module | ScopeKind::Type => {}
+            }
+        }
+        false
     }
 
     fn in_sync_comprehension(&self) -> bool {
@@ -639,6 +665,11 @@ impl SemanticSyntaxContext for Checker<'_> {
         self.semantic.current_scope().kind.is_module()
     }
 
+    fn in_function_scope(&self) -> bool {
+        let kind = &self.semantic.current_scope().kind;
+        matches!(kind, ScopeKind::Function(_) | ScopeKind::Lambda(_))
+    }
+
     fn in_notebook(&self) -> bool {
         self.source_type.is_ipynb()
     }

crates/ruff_linter/src/linter.rs

@@ -1060,4 +1060,37 @@ mod tests {
 
         Ok(())
     }
+
+    #[test_case(Path::new("yield_scope.py"); "yield_scope")]
+    fn test_yield_scope(path: &Path) -> Result<()> {
+        let snapshot = path.to_string_lossy().to_string();
+        let path = Path::new("resources/test/fixtures/syntax_errors").join(path);
+        let messages = test_contents_syntax_errors(
+            &SourceKind::Python(std::fs::read_to_string(&path)?),
+            &path,
+            &settings::LinterSettings::for_rule(Rule::YieldOutsideFunction),
+        );
+        insta::with_settings!({filters => vec![(r"\\", "/")]}, {
+            assert_messages!(snapshot, messages);
+        });
+
+        Ok(())
+    }
+
+    #[test]
+    fn test_await_scope_notebook() -> Result<()> {
+        let path = Path::new("resources/test/fixtures/syntax_errors/await_scope.ipynb");
+        let TestedNotebook {
+            messages,
+            source_notebook,
+            ..
+        } = assert_notebook_path(
+            path,
+            path,
+            &settings::LinterSettings::for_rule(Rule::YieldOutsideFunction),
+        )?;
+        assert_messages!(messages, path, source_notebook);
+
+        Ok(())
+    }
 }

crates/ruff_linter/src/rules/pyflakes/rules/yield_outside_function.rs

@@ -1,14 +1,11 @@
 use std::fmt;
 
-use ruff_diagnostics::{Diagnostic, Violation};
+use ruff_diagnostics::Violation;
 use ruff_macros::{derive_message_formats, ViolationMetadata};
-use ruff_python_ast::Expr;
+use ruff_python_parser::semantic_errors::YieldOutsideFunctionKind;
-use ruff_text_size::Ranged;
-
-use crate::checkers::ast::Checker;
 
 #[derive(Debug, PartialEq, Eq)]
-enum DeferralKeyword {
+pub(crate) enum DeferralKeyword {
     Yield,
     YieldFrom,
     Await,
@@ -24,6 +21,16 @@ impl fmt::Display for DeferralKeyword {
     }
 }
 
+impl From<YieldOutsideFunctionKind> for DeferralKeyword {
+    fn from(value: YieldOutsideFunctionKind) -> Self {
+        match value {
+            YieldOutsideFunctionKind::Yield => Self::Yield,
+            YieldOutsideFunctionKind::YieldFrom => Self::YieldFrom,
+            YieldOutsideFunctionKind::Await => Self::Await,
+        }
+    }
+}
+
 /// ## What it does
 /// Checks for `yield`, `yield from`, and `await` usages outside of functions.
 ///
@@ -50,6 +57,14 @@ pub(crate) struct YieldOutsideFunction {
     keyword: DeferralKeyword,
 }
 
+impl YieldOutsideFunction {
+    pub(crate) fn new(keyword: impl Into<DeferralKeyword>) -> Self {
+        Self {
+            keyword: keyword.into(),
+        }
+    }
+}
+
 impl Violation for YieldOutsideFunction {
     #[derive_message_formats]
     fn message(&self) -> String {
@@ -57,30 +72,3 @@ impl Violation for YieldOutsideFunction {
         format!("`{keyword}` statement outside of a function")
     }
 }
-
-/// F704
-pub(crate) fn yield_outside_function(checker: &Checker, expr: &Expr) {
-    let scope = checker.semantic().current_scope();
-    if scope.kind.is_module() || scope.kind.is_class() {
-        let keyword = match expr {
-            Expr::Yield(_) => DeferralKeyword::Yield,
-            Expr::YieldFrom(_) => DeferralKeyword::YieldFrom,
-            Expr::Await(_) => DeferralKeyword::Await,
-            _ => return,
-        };
-
-        // `await` is allowed at the top level of a Jupyter notebook.
-        // See: https://ipython.readthedocs.io/en/stable/interactive/autoawait.html.
-        if scope.kind.is_module()
-            && checker.source_type.is_ipynb()
-            && keyword == DeferralKeyword::Await
-        {
-            return;
-        }
-
-        checker.report_diagnostic(Diagnostic::new(
-            YieldOutsideFunction { keyword },
-            expr.range(),
-        ));
-    }
-}

crates/ruff_linter/src/snapshots/ruff_linter__linter__tests__await_scope_notebook.snap

@@ -0,0 +1,9 @@
+---
+source: crates/ruff_linter/src/linter.rs
+---
+await_scope.ipynb:cell 3:2:5: F704 `await` statement outside of a function
+  |
+1 | class _:
+2 |     await 1  # SyntaxError: await outside function
+  |     ^^^^^^^ F704
+  |

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

@@ -0,0 +1,117 @@
+---
+source: crates/ruff_linter/src/linter.rs
+---
+resources/test/fixtures/syntax_errors/yield_scope.py:1:1: F704 `yield` statement outside of a function
+  |
+1 | yield  # error
+  | ^^^^^ F704
+2 | yield 1  # error
+3 | yield from 1  # error
+  |
+
+resources/test/fixtures/syntax_errors/yield_scope.py:2:1: F704 `yield` statement outside of a function
+  |
+1 | yield  # error
+2 | yield 1  # error
+  | ^^^^^^^ F704
+3 | yield from 1  # error
+4 | await 1  # error
+  |
+
+resources/test/fixtures/syntax_errors/yield_scope.py:3:1: F704 `yield from` statement outside of a function
+  |
+1 | yield  # error
+2 | yield 1  # error
+3 | yield from 1  # error
+  | ^^^^^^^^^^^^ F704
+4 | await 1  # error
+5 | [(yield x) for x in range(3)]  # error
+  |
+
+resources/test/fixtures/syntax_errors/yield_scope.py:4:1: F704 `await` statement outside of a function
+  |
+2 | yield 1  # error
+3 | yield from 1  # error
+4 | await 1  # error
+  | ^^^^^^^ F704
+5 | [(yield x) for x in range(3)]  # error
+  |
+
+resources/test/fixtures/syntax_errors/yield_scope.py:5:3: F704 `yield` statement outside of a function
+  |
+3 | yield from 1  # error
+4 | await 1  # error
+5 | [(yield x) for x in range(3)]  # error
+  |   ^^^^^^^ F704
+  |
+
+resources/test/fixtures/syntax_errors/yield_scope.py:23:9: F704 `yield` statement outside of a function
+   |
+21 | def outer():
+22 |     class C:
+23 |         yield 1  # error
+   |         ^^^^^^^ F704
+24 |
+25 |     [(yield 1) for x in range(3)]  # error
+   |
+
+resources/test/fixtures/syntax_errors/yield_scope.py:25:7: F704 `yield` statement outside of a function
+   |
+23 |         yield 1  # error
+24 |
+25 |     [(yield 1) for x in range(3)]  # error
+   |       ^^^^^^^ F704
+26 |     ((yield 1) for x in range(3))  # error
+27 |     {(yield 1) for x in range(3)}  # error
+   |
+
+resources/test/fixtures/syntax_errors/yield_scope.py:26:7: F704 `yield` statement outside of a function
+   |
+25 |     [(yield 1) for x in range(3)]  # error
+26 |     ((yield 1) for x in range(3))  # error
+   |       ^^^^^^^ F704
+27 |     {(yield 1) for x in range(3)}  # error
+28 |     {(yield 1): 0 for x in range(3)}  # error
+   |
+
+resources/test/fixtures/syntax_errors/yield_scope.py:27:7: F704 `yield` statement outside of a function
+   |
+25 |     [(yield 1) for x in range(3)]  # error
+26 |     ((yield 1) for x in range(3))  # error
+27 |     {(yield 1) for x in range(3)}  # error
+   |       ^^^^^^^ F704
+28 |     {(yield 1): 0 for x in range(3)}  # error
+29 |     {0: (yield 1) for x in range(3)}  # error
+   |
+
+resources/test/fixtures/syntax_errors/yield_scope.py:28:7: F704 `yield` statement outside of a function
+   |
+26 |     ((yield 1) for x in range(3))  # error
+27 |     {(yield 1) for x in range(3)}  # error
+28 |     {(yield 1): 0 for x in range(3)}  # error
+   |       ^^^^^^^ F704
+29 |     {0: (yield 1) for x in range(3)}  # error
+   |
+
+resources/test/fixtures/syntax_errors/yield_scope.py:29:10: F704 `yield` statement outside of a function
+   |
+27 |     {(yield 1) for x in range(3)}  # error
+28 |     {(yield 1): 0 for x in range(3)}  # error
+29 |     {0: (yield 1) for x in range(3)}  # error
+   |          ^^^^^^^ F704
+   |
+
+resources/test/fixtures/syntax_errors/yield_scope.py:36:10: F704 `await` statement outside of a function
+   |
+35 |     class C:
+36 |         [await x for x in range(3)]  # error, classes break async scope
+   |          ^^^^^^^ F704
+37 |
+38 |     lambda x: await x  # okay for now, lambda breaks _async_ scope but is a function
+   |
+
+resources/test/fixtures/syntax_errors/yield_scope.py:41:1: F704 `await` statement outside of a function
+   |
+41 | await 1  # error
+   | ^^^^^^^ F704
+   |

crates/ruff_python_parser/src/semantic_errors.rs

@@ -587,17 +587,53 @@ impl SemanticSyntaxChecker {
                     }
                 }
             }
-            Expr::Yield(ast::ExprYield {
+            Expr::Yield(ast::ExprYield { value, .. }) => {
-                value: Some(value), ..
+                if let Some(value) = value {
-            }) => {
                     // test_err single_star_yield
                     // def f(): yield *x
                     Self::invalid_star_expression(value, ctx);
                 }
+                Self::yield_outside_function(ctx, expr, YieldOutsideFunctionKind::Yield);
+            }
+            Expr::YieldFrom(_) => {
+                Self::yield_outside_function(ctx, expr, YieldOutsideFunctionKind::YieldFrom);
+            }
+            Expr::Await(_) => {
+                Self::yield_outside_function(ctx, expr, YieldOutsideFunctionKind::Await);
+            }
             _ => {}
         }
     }
 
+    /// F704
+    fn yield_outside_function<Ctx: SemanticSyntaxContext>(
+        ctx: &Ctx,
+        expr: &Expr,
+        kind: YieldOutsideFunctionKind,
+    ) {
+        // We are intentionally not inspecting the async status of the scope for now to mimic F704.
+        // await-outside-async is PLE1142 instead, so we'll end up emitting both syntax errors for
+        // cases that trigger F704
+        if kind.is_await() {
+            if ctx.in_await_allowed_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;
+            }
+        } else if ctx.in_function_scope() {
+            return;
+        }
+
+        Self::add_error(
+            ctx,
+            SemanticSyntaxErrorKind::YieldOutsideFunction(kind),
+            expr.range(),
+        );
+    }
+
     /// Add a [`SyntaxErrorKind::ReboundComprehensionVariable`] if `expr` rebinds an iteration
     /// variable in `generators`.
     fn check_generator_expr<Ctx: SemanticSyntaxContext>(
@@ -758,6 +794,9 @@ impl Display for SemanticSyntaxError {
                                 function on Python {python_version} (syntax was added in 3.11)",
                 )
             }
+            SemanticSyntaxErrorKind::YieldOutsideFunction(kind) => {
+                write!(f, "`{kind}` statement outside of a function")
+            }
         }
     }
 }
@@ -1013,6 +1052,69 @@ pub enum SemanticSyntaxErrorKind {
     ///
     /// [BPO 33346]: https://github.com/python/cpython/issues/77527
     AsyncComprehensionOutsideAsyncFunction(PythonVersion),
+
+    /// Represents the use of `yield`, `yield from`, or `await` outside of a function scope.
+    ///
+    ///
+    /// ## Examples
+    ///
+    /// `yield` and `yield from` are only allowed if the immediately-enclosing scope is a function
+    /// or lambda and not allowed otherwise:
+    ///
+    /// ```python
+    /// yield 1  # error
+    ///
+    /// def f():
+    ///     [(yield 1) for x in y]  # error
+    /// ```
+    ///
+    /// `await` is additionally allowed in comprehensions, if the comprehension itself is in a
+    /// function scope:
+    ///
+    /// ```python
+    /// await 1  # error
+    ///
+    /// async def f():
+    ///     await 1  # okay
+    ///     [await 1 for x in y]  # also okay
+    /// ```
+    ///
+    /// This last case _is_ an error, but it has to do with the lambda not being an async function.
+    /// For the sake of this error kind, this is okay.
+    ///
+    /// ## References
+    ///
+    /// See [PEP 255] for details on `yield`, [PEP 380] for the extension to `yield from`, [PEP 492]
+    /// for async-await syntax, and [PEP 530] for async comprehensions.
+    ///
+    /// [PEP 255]: https://peps.python.org/pep-0255/
+    /// [PEP 380]: https://peps.python.org/pep-0380/
+    /// [PEP 492]: https://peps.python.org/pep-0492/
+    /// [PEP 530]: https://peps.python.org/pep-0530/
+    YieldOutsideFunction(YieldOutsideFunctionKind),
+}
+
+#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
+pub enum YieldOutsideFunctionKind {
+    Yield,
+    YieldFrom,
+    Await,
+}
+
+impl YieldOutsideFunctionKind {
+    pub fn is_await(&self) -> bool {
+        matches!(self, Self::Await)
+    }
+}
+
+impl Display for YieldOutsideFunctionKind {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        f.write_str(match self {
+            YieldOutsideFunctionKind::Yield => "yield",
+            YieldOutsideFunctionKind::YieldFrom => "yield from",
+            YieldOutsideFunctionKind::Await => "await",
+        })
+    }
 }
 
 #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
@@ -1326,6 +1428,40 @@ where
     }
 }
 
+/// Information needed from a parent visitor to emit semantic syntax errors.
+///
+/// Note that the `in_*_scope` methods should refer to the immediately-enclosing scope. For example,
+/// `in_function_scope` should return true for this case:
+///
+/// ```python
+/// def f():
+///     x  # here
+/// ```
+///
+/// but not for this case:
+///
+/// ```python
+/// def f():
+///     class C:
+///         x  # here
+/// ```
+///
+/// In contrast, the `in_*_context` methods should traverse parent scopes. For example,
+/// `in_function_context` should return true for this case:
+///
+/// ```python
+/// def f():
+///     [x  # here
+///         for x in range(3)]
+/// ```
+///
+/// but not here:
+///
+/// ```python
+/// def f():
+///     class C:
+///         x  # here, classes break function scopes
+/// ```
 pub trait SemanticSyntaxContext {
     /// Returns `true` if a module's docstring boundary has been passed.
     fn seen_docstring_boundary(&self) -> bool;
@@ -1345,6 +1481,29 @@ pub trait SemanticSyntaxContext {
     /// Returns `true` if the visitor is currently in an async context, i.e. an async function.
     fn in_async_context(&self) -> bool;
 
+    /// Returns `true` if the visitor is currently in a context where the `await` keyword is
+    /// allowed.
+    ///
+    /// Note that this is method is primarily used to report `YieldOutsideFunction` errors for
+    /// `await` outside function scopes, irrespective of their async status. As such, this differs
+    /// from `in_async_context` in two ways:
+    ///
+    /// 1. `await` is allowed in a lambda, despite it not being async
+    /// 2. `await` is allowed in any function, regardless of its async status
+    ///
+    /// In short, only nested class definitions should cause this method to return `false`, for
+    /// example:
+    ///
+    /// ```python
+    /// def f():
+    ///     await 1  # okay, in a function
+    ///     class C:
+    ///         await 1  # error
+    /// ```
+    ///
+    /// See the trait-level documentation for more details.
+    fn in_await_allowed_context(&self) -> bool;
+
     /// Returns `true` if the visitor is currently inside of a synchronous comprehension.
     ///
     /// This method is necessary because `in_async_context` only checks for the nearest, enclosing
@@ -1356,6 +1515,9 @@ pub trait SemanticSyntaxContext {
     /// Returns `true` if the visitor is at the top-level module scope.
     fn in_module_scope(&self) -> bool;
 
+    /// Returns `true` if the visitor is in a function scope.
+    fn in_function_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

@@ -464,6 +464,7 @@ enum Scope {
     Module,
     Function { is_async: bool },
     Comprehension { is_async: bool },
+    Class,
 }
 
 struct SemanticSyntaxCheckerVisitor<'a> {
@@ -546,20 +547,46 @@ impl SemanticSyntaxContext for SemanticSyntaxCheckerVisitor<'_> {
     }
 
     fn in_module_scope(&self) -> bool {
-        self.scopes
+        true
-            .last()
+    }
-            .is_some_and(|scope| matches!(scope, Scope::Module))
+
+    fn in_function_scope(&self) -> bool {
+        true
     }
 
     fn in_notebook(&self) -> bool {
         false
     }
+
+    fn in_await_allowed_context(&self) -> bool {
+        true
+    }
 }
 
 impl Visitor<'_> for SemanticSyntaxCheckerVisitor<'_> {
     fn visit_stmt(&mut self, stmt: &ast::Stmt) {
         self.with_semantic_checker(|semantic, context| semantic.visit_stmt(stmt, context));
         match stmt {
+            ast::Stmt::ClassDef(ast::StmtClassDef {
+                arguments,
+                body,
+                decorator_list,
+                type_params,
+                ..
+            }) => {
+                for decorator in decorator_list {
+                    self.visit_decorator(decorator);
+                }
+                if let Some(type_params) = type_params {
+                    self.visit_type_params(type_params);
+                }
+                if let Some(arguments) = arguments {
+                    self.visit_arguments(arguments);
+                }
+                self.scopes.push(Scope::Class);
+                self.visit_body(body);
+                self.scopes.pop().unwrap();
+            }
             ast::Stmt::FunctionDef(ast::StmtFunctionDef { is_async, .. }) => {
                 self.scopes.push(Scope::Function {
                     is_async: *is_async,
@@ -581,13 +608,38 @@ impl Visitor<'_> for SemanticSyntaxCheckerVisitor<'_> {
                 ast::visitor::walk_expr(self, expr);
                 self.scopes.pop().unwrap();
             }
-            ast::Expr::ListComp(ast::ExprListComp { generators, .. })
+            ast::Expr::ListComp(ast::ExprListComp {
-            | ast::Expr::SetComp(ast::ExprSetComp { generators, .. })
+                elt, generators, ..
-            | ast::Expr::DictComp(ast::ExprDictComp { generators, .. }) => {
+            })
+            | ast::Expr::SetComp(ast::ExprSetComp {
+                elt, generators, ..
+            })
+            | ast::Expr::Generator(ast::ExprGenerator {
+                elt, generators, ..
+            }) => {
+                for comprehension in generators {
+                    self.visit_comprehension(comprehension);
+                }
                 self.scopes.push(Scope::Comprehension {
                     is_async: generators.iter().any(|gen| gen.is_async),
                 });
-                ast::visitor::walk_expr(self, expr);
+                self.visit_expr(elt);
+                self.scopes.pop().unwrap();
+            }
+            ast::Expr::DictComp(ast::ExprDictComp {
+                key,
+                value,
+                generators,
+                ..
+            }) => {
+                for comprehension in generators {
+                    self.visit_comprehension(comprehension);
+                }
+                self.scopes.push(Scope::Comprehension {
+                    is_async: generators.iter().any(|gen| gen.is_async),
+                });
+                self.visit_expr(key);
+                self.visit_expr(value);
                 self.scopes.pop().unwrap();
             }
             _ => {