Remove dedicated ScopeKind structs in favor of AST nodes (#4648)

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

@@ -31,7 +31,7 @@ use ruff_python_semantic::context::ExecutionContext;
 use ruff_python_semantic::definition::{ContextualizedDefinition, Module, ModuleKind};
 use ruff_python_semantic::model::{ResolvedReference, SemanticModel, SemanticModelFlags};
 use ruff_python_semantic::node::NodeId;
-use ruff_python_semantic::scope::{ClassDef, FunctionDef, Lambda, Scope, ScopeId, ScopeKind};
+use ruff_python_semantic::scope::{Scope, ScopeId, ScopeKind};
 use ruff_python_stdlib::builtins::{BUILTINS, MAGIC_GLOBALS};
 use ruff_python_stdlib::path::is_python_stub_file;
 
@@ -1683,7 +1683,7 @@ where
                         if !self
                             .semantic_model
                             .scopes()
-                            .any(|scope| scope.kind.is_function())
+                            .any(|scope| scope.kind.is_any_function())
                         {
                             if self.enabled(Rule::UnprefixedTypeParam) {
                                 flake8_pyi::rules::prefix_type_params(self, value, targets);
@@ -1732,7 +1732,7 @@ where
                             if !self
                                 .semantic_model
                                 .scopes()
-                                .any(|scope| scope.kind.is_function())
+                                .any(|scope| scope.kind.is_any_function())
                             {
                                 flake8_pyi::rules::annotated_assignment_default_in_stub(
                                     self, target, value, annotation,
@@ -1886,9 +1886,25 @@ where
                     },
                 );
 
-                // If any global bindings don't already exist in the global scope, add it.
+                let definition = docstrings::extraction::extract_definition(
+                    ExtractionTarget::Function,
+                    stmt,
+                    self.semantic_model.definition_id,
+                    &self.semantic_model.definitions,
+                );
+                self.semantic_model.push_definition(definition);
+
+                self.semantic_model.push_scope(match &stmt {
+                    Stmt::FunctionDef(stmt) => ScopeKind::Function(stmt),
+                    Stmt::AsyncFunctionDef(stmt) => ScopeKind::AsyncFunction(stmt),
+                    _ => unreachable!("Expected Stmt::FunctionDef | Stmt::AsyncFunctionDef"),
+                });
+
+                self.deferred.functions.push(self.semantic_model.snapshot());
+
+                // If any global bindings don't already exist in the global scope, add them.
                 let globals = helpers::extract_globals(body);
-                for (name, stmt) in helpers::extract_globals(body) {
+                for (name, range) in globals {
                     if self
                         .semantic_model
                         .global_scope()
@@ -1901,7 +1917,7 @@ where
                     {
                         let id = self.semantic_model.bindings.push(Binding {
                             kind: BindingKind::Assignment,
-                            range: stmt.range(),
+                            range,
                             references: Vec::new(),
                             source: self.semantic_model.stmt_id,
                             context: self.semantic_model.execution_context(),
@@ -1910,36 +1926,19 @@ where
                         });
                         self.semantic_model.global_scope_mut().add(name, id);
                     }
+
+                    self.semantic_model.scope_mut().add_global(name, range);
                 }
-
-                let definition = docstrings::extraction::extract_definition(
-                    ExtractionTarget::Function,
-                    stmt,
-                    self.semantic_model.definition_id,
-                    &self.semantic_model.definitions,
-                );
-                self.semantic_model.push_definition(definition);
-
-                self.semantic_model
-                    .push_scope(ScopeKind::Function(FunctionDef {
-                        name,
-                        body,
-                        args,
-                        decorator_list,
-                        async_: matches!(stmt, Stmt::AsyncFunctionDef(_)),
-                        globals,
-                    }));
-
-                self.deferred.functions.push(self.semantic_model.snapshot());
             }
-            Stmt::ClassDef(ast::StmtClassDef {
+            Stmt::ClassDef(
-                body,
+                class_def @ ast::StmtClassDef {
-                name,
+                    body,
-                bases,
+                    bases,
-                keywords,
+                    keywords,
-                decorator_list,
+                    decorator_list,
-                range: _,
+                    ..
-            }) => {
+                },
+            ) => {
                 for expr in bases {
                     self.visit_expr(expr);
                 }
@@ -1950,9 +1949,18 @@ where
                     self.visit_expr(expr);
                 }
 
-                // If any global bindings don't already exist in the global scope, add it.
+                let definition = docstrings::extraction::extract_definition(
-                let globals = helpers::extract_globals(body);
+                    ExtractionTarget::Class,
-                for (name, stmt) in &globals {
+                    stmt,
+                    self.semantic_model.definition_id,
+                    &self.semantic_model.definitions,
+                );
+                self.semantic_model.push_definition(definition);
+
+                self.semantic_model.push_scope(ScopeKind::Class(class_def));
+
+                // If any global bindings don't already exist in the global scope, add them.
+                for (name, range) in helpers::extract_globals(body) {
                     if self
                         .semantic_model
                         .global_scope()
@@ -1965,7 +1973,7 @@ where
                     {
                         let id = self.semantic_model.bindings.push(Binding {
                             kind: BindingKind::Assignment,
-                            range: stmt.range(),
+                            range,
                             references: Vec::new(),
                             source: self.semantic_model.stmt_id,
                             context: self.semantic_model.execution_context(),
@@ -1974,24 +1982,9 @@ where
                         });
                         self.semantic_model.global_scope_mut().add(name, id);
                     }
+                    self.semantic_model.scope_mut().add_global(name, range);
                 }
 
-                let definition = docstrings::extraction::extract_definition(
-                    ExtractionTarget::Class,
-                    stmt,
-                    self.semantic_model.definition_id,
-                    &self.semantic_model.definitions,
-                );
-                self.semantic_model.push_definition(definition);
-
-                self.semantic_model.push_scope(ScopeKind::Class(ClassDef {
-                    name,
-                    bases,
-                    keywords,
-                    decorator_list,
-                    globals,
-                }));
-
                 self.visit_body(body);
             }
             Stmt::Try(ast::StmtTry {
@@ -2054,7 +2047,7 @@ where
                 // available at runtime.
                 // See: https://docs.python.org/3/reference/simple_stmts.html#annotated-assignment-statements
                 let runtime_annotation = if self.semantic_model.future_annotations() {
-                    if matches!(self.semantic_model.scope().kind, ScopeKind::Class(..)) {
+                    if self.semantic_model.scope().kind.is_class() {
                         let baseclasses = &self
                             .settings
                             .flake8_type_checking
@@ -2074,7 +2067,7 @@ where
                 } else {
                     matches!(
                         self.semantic_model.scope().kind,
-                        ScopeKind::Class(..) | ScopeKind::Module
+                        ScopeKind::Class(_) | ScopeKind::Module
                     )
                 };
 
@@ -3419,7 +3412,7 @@ where
             Expr::Lambda(
                 lambda @ ast::ExprLambda {
                     args,
-                    body,
+                    body: _,
                     range: _,
                 },
             ) => {
@@ -3434,8 +3427,7 @@ where
                 for expr in &args.defaults {
                     self.visit_expr(expr);
                 }
-                self.semantic_model
+                self.semantic_model.push_scope(ScopeKind::Lambda(lambda));
-                    .push_scope(ScopeKind::Lambda(Lambda { args, body }));
             }
             Expr::IfExp(ast::ExprIfExp {
                 test,
@@ -4514,7 +4506,7 @@ impl<'a> Checker<'a> {
         }
 
         if self.enabled(Rule::NonLowercaseVariableInFunction) {
-            if matches!(self.semantic_model.scope().kind, ScopeKind::Function(..)) {
+            if self.semantic_model.scope().kind.is_any_function() {
                 // Ignore globals.
                 if !self
                     .semantic_model
@@ -4530,13 +4522,11 @@ impl<'a> Checker<'a> {
         }
 
         if self.enabled(Rule::MixedCaseVariableInClassScope) {
-            if let ScopeKind::Class(class) = &self.semantic_model.scope().kind {
+            if let ScopeKind::Class(ast::StmtClassDef { bases, .. }) =
+                &self.semantic_model.scope().kind
+            {
                 pep8_naming::rules::mixed_case_variable_in_class_scope(
-                    self,
+                    self, expr, parent, id, bases,
-                    expr,
-                    parent,
-                    id,
-                    class.bases,
                 );
             }
         }
@@ -5058,7 +5048,7 @@ impl<'a> Checker<'a> {
             }
 
             // Imports in classes are public members.
-            if matches!(scope.kind, ScopeKind::Class(..)) {
+            if scope.kind.is_class() {
                 continue;
             }
 

crates/ruff/src/rules/flake8_async/helpers.rs

@@ -1,18 +0,0 @@
-use ruff_python_semantic::{
-    model::SemanticModel,
-    scope::{FunctionDef, ScopeKind},
-};
-
-/// Return `true` if the [`SemanticModel`] is inside an async function definition.
-pub(crate) fn in_async_function(model: &SemanticModel) -> bool {
-    model
-        .scopes()
-        .find_map(|scope| {
-            if let ScopeKind::Function(FunctionDef { async_, .. }) = &scope.kind {
-                Some(*async_)
-            } else {
-                None
-            }
-        })
-        .unwrap_or(false)
-}

crates/ruff/src/rules/flake8_async/mod.rs

@@ -1,5 +1,4 @@
 //! Rules from [flake8-async](https://pypi.org/project/flake8-async/).
-mod helpers;
 pub(crate) mod rules;
 
 #[cfg(test)]

crates/ruff/src/rules/flake8_async/rules/blocking_http_call.rs

@@ -6,8 +6,6 @@ use ruff_macros::{derive_message_formats, violation};
 
 use crate::checkers::ast::Checker;
 
-use super::super::helpers::in_async_function;
-
 /// ## What it does
 /// Checks that async functions do not contain blocking HTTP calls.
 ///
@@ -66,7 +64,7 @@ const BLOCKING_HTTP_CALLS: &[&[&str]] = &[
 
 /// ASYNC100
 pub(crate) fn blocking_http_call(checker: &mut Checker, expr: &Expr) {
-    if in_async_function(checker.semantic_model()) {
+    if checker.semantic_model().in_async_context() {
         if let Expr::Call(ast::ExprCall { func, .. }) = expr {
             let call_path = checker.semantic_model().resolve_call_path(func);
             let is_blocking =

crates/ruff/src/rules/flake8_async/rules/blocking_os_call.rs

@@ -6,8 +6,6 @@ use ruff_macros::{derive_message_formats, violation};
 
 use crate::checkers::ast::Checker;
 
-use super::super::helpers::in_async_function;
-
 /// ## What it does
 /// Checks that async functions do not contain calls to blocking synchronous
 /// process calls via the `os` module.
@@ -58,7 +56,7 @@ const UNSAFE_OS_METHODS: &[&[&str]] = &[
 
 /// ASYNC102
 pub(crate) fn blocking_os_call(checker: &mut Checker, expr: &Expr) {
-    if in_async_function(checker.semantic_model()) {
+    if checker.semantic_model().in_async_context() {
         if let Expr::Call(ast::ExprCall { func, .. }) = expr {
             let is_unsafe_os_method = checker
                 .semantic_model()

crates/ruff/src/rules/flake8_async/rules/open_sleep_or_subprocess_call.rs

@@ -6,8 +6,6 @@ use ruff_macros::{derive_message_formats, violation};
 
 use crate::checkers::ast::Checker;
 
-use super::super::helpers::in_async_function;
-
 /// ## What it does
 /// Checks that async functions do not contain calls to `open`, `time.sleep`,
 /// or `subprocess` methods.
@@ -61,7 +59,7 @@ const OPEN_SLEEP_OR_SUBPROCESS_CALL: &[&[&str]] = &[
 
 /// ASYNC101
 pub(crate) fn open_sleep_or_subprocess_call(checker: &mut Checker, expr: &Expr) {
-    if in_async_function(checker.semantic_model()) {
+    if checker.semantic_model().in_async_context() {
         if let Expr::Call(ast::ExprCall { func, .. }) = expr {
             let is_open_sleep_or_subprocess_call = checker
                 .semantic_model()

crates/ruff/src/rules/flake8_bugbear/rules/cached_instance_method.rs

@@ -3,7 +3,6 @@ use rustpython_parser::ast::{self, Expr, Ranged};
 use ruff_diagnostics::{Diagnostic, Violation};
 use ruff_macros::{derive_message_formats, violation};
 use ruff_python_semantic::model::SemanticModel;
-use ruff_python_semantic::scope::ScopeKind;
 
 use crate::checkers::ast::Checker;
 
@@ -28,7 +27,7 @@ fn is_cache_func(model: &SemanticModel, expr: &Expr) -> bool {
 
 /// B019
 pub(crate) fn cached_instance_method(checker: &mut Checker, decorator_list: &[Expr]) {
-    if !matches!(checker.semantic_model().scope().kind, ScopeKind::Class(_)) {
+    if !checker.semantic_model().scope().kind.is_class() {
         return;
     }
     for decorator in decorator_list {

crates/ruff/src/rules/flake8_pyi/rules/simple_defaults.rs

@@ -4,7 +4,7 @@ use ruff_diagnostics::{AlwaysAutofixableViolation, Diagnostic, Edit, Fix, Violat
 use ruff_macros::{derive_message_formats, violation};
 use ruff_python_ast::source_code::Locator;
 use ruff_python_semantic::model::SemanticModel;
-use ruff_python_semantic::scope::{ClassDef, ScopeKind};
+use ruff_python_semantic::scope::ScopeKind;
 
 use crate::checkers::ast::Checker;
 use crate::registry::AsRule;
@@ -558,7 +558,8 @@ pub(crate) fn unannotated_assignment_in_stub(
         return;
     }
 
-    if let ScopeKind::Class(ClassDef { bases, .. }) = checker.semantic_model().scope().kind {
+    if let ScopeKind::Class(ast::StmtClassDef { bases, .. }) = checker.semantic_model().scope().kind
+    {
         if is_enum(checker.semantic_model(), bases) {
             return;
         }

crates/ruff/src/rules/flake8_self/rules/private_member_access.rs

@@ -99,19 +99,19 @@ pub(crate) fn private_member_access(checker: &mut Checker, expr: &Expr) {
                         .iter()
                         .rev()
                         .find_map(|scope| match &scope.kind {
-                            ScopeKind::Class(class_def) => Some(class_def),
+                            ScopeKind::Class(ast::StmtClassDef { name, .. }) => Some(name),
                             _ => None,
                         })
-                        .map_or(false, |class_def| {
+                        .map_or(false, |name| {
-                            if call_path.as_slice() == [class_def.name] {
+                            if call_path.as_slice() == [name.as_str()] {
-                                checker
+                                checker.semantic_model().find_binding(name).map_or(
-                                    .semantic_model()
+                                    false,
-                                    .find_binding(class_def.name)
+                                    |binding| {
-                                    .map_or(false, |binding| {
                                         // TODO(charlie): Could the name ever be bound to a
                                         // _different_ class here?
                                         binding.kind.is_class_definition()
-                                    })
+                                    },
+                                )
                             } else {
                                 false
                             }

crates/ruff/src/rules/flake8_simplify/rules/ast_unary_op.rs

@@ -98,8 +98,11 @@ pub(crate) fn negation_with_equal_op(
     }
 
     // Avoid flagging issues in dunder implementations.
-    if let ScopeKind::Function(def) = &checker.semantic_model().scope().kind {
+    if let ScopeKind::Function(ast::StmtFunctionDef { name, .. })
-        if DUNDER_METHODS.contains(&def.name) {
+    | ScopeKind::AsyncFunction(ast::StmtAsyncFunctionDef { name, .. }) =
+        &checker.semantic_model().scope().kind
+    {
+        if DUNDER_METHODS.contains(&name.as_str()) {
             return;
         }
     }
@@ -148,8 +151,11 @@ pub(crate) fn negation_with_not_equal_op(
     }
 
     // Avoid flagging issues in dunder implementations.
-    if let ScopeKind::Function(def) = &checker.semantic_model().scope().kind {
+    if let ScopeKind::Function(ast::StmtFunctionDef { name, .. })
-        if DUNDER_METHODS.contains(&def.name) {
+    | ScopeKind::AsyncFunction(ast::StmtAsyncFunctionDef { name, .. }) =
+        &checker.semantic_model().scope().kind
+    {
+        if DUNDER_METHODS.contains(&name.as_str()) {
             return;
         }
     }

crates/ruff/src/rules/flake8_type_checking/helpers.rs

@@ -83,8 +83,8 @@ pub(crate) fn runtime_evaluated(
 }
 
 fn runtime_evaluated_base_class(semantic_model: &SemanticModel, base_classes: &[String]) -> bool {
-    if let ScopeKind::Class(class_def) = &semantic_model.scope().kind {
+    if let ScopeKind::Class(ast::StmtClassDef { bases, .. }) = &semantic_model.scope().kind {
-        for base in class_def.bases.iter() {
+        for base in bases.iter() {
             if let Some(call_path) = semantic_model.resolve_call_path(base) {
                 if base_classes
                     .iter()
@@ -99,8 +99,9 @@ fn runtime_evaluated_base_class(semantic_model: &SemanticModel, base_classes: &[
 }
 
 fn runtime_evaluated_decorators(semantic_model: &SemanticModel, decorators: &[String]) -> bool {
-    if let ScopeKind::Class(class_def) = &semantic_model.scope().kind {
+    if let ScopeKind::Class(ast::StmtClassDef { decorator_list, .. }) = &semantic_model.scope().kind
-        for decorator in class_def.decorator_list.iter() {
+    {
+        for decorator in decorator_list.iter() {
             if let Some(call_path) = semantic_model.resolve_call_path(map_callable(decorator)) {
                 if decorators
                     .iter()

crates/ruff/src/rules/flake8_unused_arguments/rules/unused_arguments.rs

@@ -1,6 +1,7 @@
 use std::iter;
 
 use regex::Regex;
+use rustpython_parser::ast;
 use rustpython_parser::ast::{Arg, Arguments};
 
 use ruff_diagnostics::DiagnosticKind;
@@ -10,7 +11,7 @@ use ruff_python_semantic::analyze::function_type;
 use ruff_python_semantic::analyze::function_type::FunctionType;
 use ruff_python_semantic::analyze::visibility;
 use ruff_python_semantic::binding::Bindings;
-use ruff_python_semantic::scope::{FunctionDef, Lambda, Scope, ScopeKind};
+use ruff_python_semantic::scope::{Scope, ScopeKind};
 
 use super::super::helpers;
 
@@ -317,7 +318,14 @@ pub(crate) fn unused_arguments(
     bindings: &Bindings,
 ) -> Vec<Diagnostic> {
     match &scope.kind {
-        ScopeKind::Function(FunctionDef {
+        ScopeKind::Function(ast::StmtFunctionDef {
+            name,
+            args,
+            body,
+            decorator_list,
+            ..
+        })
+        | ScopeKind::AsyncFunction(ast::StmtAsyncFunctionDef {
             name,
             args,
             body,
@@ -431,7 +439,7 @@ pub(crate) fn unused_arguments(
                 }
             }
         }
-        ScopeKind::Lambda(Lambda { args, .. }) => {
+        ScopeKind::Lambda(ast::ExprLambda { args, .. }) => {
             if checker.enabled(Argumentable::Lambda.rule_code()) {
                 function(
                     Argumentable::Lambda,

crates/ruff/src/rules/pycodestyle/rules/lambda_assignment.rs

@@ -8,7 +8,6 @@ use ruff_python_ast::newlines::StrExt;
 use ruff_python_ast::source_code::Generator;
 use ruff_python_ast::whitespace::leading_space;
 use ruff_python_semantic::model::SemanticModel;
-use ruff_python_semantic::scope::ScopeKind;
 
 use crate::checkers::ast::Checker;
 use crate::registry::AsRule;
@@ -66,14 +65,6 @@ pub(crate) fn lambda_assignment(
 ) {
     if let Expr::Name(ast::ExprName { id, .. }) = target {
         if let Expr::Lambda(ast::ExprLambda { args, body, .. }) = value {
-            // If the assignment is in a class body, it might not be safe
-            // to replace it because the assignment might be
-            // carrying a type annotation that will be used by some
-            // package like dataclasses, which wouldn't consider the
-            // rewritten function definition to be equivalent.
-            // See https://github.com/charliermarsh/ruff/issues/3046
-            let fixable = !matches!(checker.semantic_model().scope().kind, ScopeKind::Class(_));
-
             let mut diagnostic = Diagnostic::new(
                 LambdaAssignment {
                     name: id.to_string(),
@@ -81,8 +72,14 @@ pub(crate) fn lambda_assignment(
                 stmt.range(),
             );
 
+            // If the assignment is in a class body, it might not be safe
+            // to replace it because the assignment might be
+            // carrying a type annotation that will be used by some
+            // package like dataclasses, which wouldn't consider the
+            // rewritten function definition to be equivalent.
+            // See https://github.com/charliermarsh/ruff/issues/3046
             if checker.patch(diagnostic.kind.rule())
-                && fixable
+                && !checker.semantic_model().scope().kind.is_class()
                 && !has_leading_content(stmt, checker.locator)
                 && !has_trailing_content(stmt, checker.locator)
             {

crates/ruff/src/rules/pyflakes/rules/undefined_local.rs

@@ -22,7 +22,7 @@ impl Violation for UndefinedLocal {
 pub(crate) fn undefined_local(checker: &mut Checker, name: &str) {
     // If the name hasn't already been defined in the current scope...
     let current = checker.semantic_model().scope();
-    if !current.kind.is_function() || current.defines(name) {
+    if !current.kind.is_any_function() || current.defines(name) {
         return;
     }
 
@@ -36,7 +36,7 @@ pub(crate) fn undefined_local(checker: &mut Checker, name: &str) {
         .scopes
         .ancestors(parent)
         .find_map(|scope| {
-            if !(scope.kind.is_function() || scope.kind.is_module()) {
+            if !(scope.kind.is_any_function() || scope.kind.is_module()) {
                 return None;
             }
 

crates/ruff/src/rules/pylint/helpers.rs

@@ -1,17 +1,25 @@
 use ruff_python_semantic::analyze::function_type;
 use ruff_python_semantic::analyze::function_type::FunctionType;
 use ruff_python_semantic::model::SemanticModel;
-use ruff_python_semantic::scope::{FunctionDef, ScopeKind};
+use ruff_python_semantic::scope::ScopeKind;
+use rustpython_parser::ast;
 
 use crate::settings::Settings;
 
 pub(crate) fn in_dunder_init(model: &SemanticModel, settings: &Settings) -> bool {
     let scope = model.scope();
-    let ScopeKind::Function(FunctionDef {
+    let (
-        name,
+        ScopeKind::Function(ast::StmtFunctionDef {
-        decorator_list,
+            name,
+            decorator_list,
         ..
-    }): ScopeKind = scope.kind else {
+        }) |
+        ScopeKind::AsyncFunction(ast::StmtAsyncFunctionDef {
+            name,
+            decorator_list,
+            ..
+        })
+    ) = scope.kind else {
         return false;
     };
     if name != "__init__" {

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

@@ -2,7 +2,6 @@ use rustpython_parser::ast::{Expr, Ranged};
 
 use ruff_diagnostics::{Diagnostic, Violation};
 use ruff_macros::{derive_message_formats, violation};
-use ruff_python_semantic::scope::{FunctionDef, ScopeKind};
 
 use crate::checkers::ast::Checker;
 
@@ -45,18 +44,7 @@ impl Violation for AwaitOutsideAsync {
 
 /// PLE1142
 pub(crate) fn await_outside_async(checker: &mut Checker, expr: &Expr) {
-    if !checker
+    if !checker.semantic_model().in_async_context() {
-        .semantic_model()
-        .scopes()
-        .find_map(|scope| {
-            if let ScopeKind::Function(FunctionDef { async_, .. }) = &scope.kind {
-                Some(*async_)
-            } else {
-                None
-            }
-        })
-        .unwrap_or(true)
-    {
         checker
             .diagnostics
             .push(Diagnostic::new(AwaitOutsideAsync, expr.range()));

crates/ruff/src/rules/pylint/rules/load_before_global_declaration.rs

@@ -3,7 +3,6 @@ use rustpython_parser::ast::{Expr, Ranged};
 use ruff_diagnostics::{Diagnostic, Violation};
 use ruff_macros::{derive_message_formats, violation};
 use ruff_python_ast::source_code::OneIndexed;
-use ruff_python_semantic::scope::ScopeKind;
 
 use crate::checkers::ast::Checker;
 
@@ -55,12 +54,8 @@ impl Violation for LoadBeforeGlobalDeclaration {
 }
 /// PLE0118
 pub(crate) fn load_before_global_declaration(checker: &mut Checker, name: &str, expr: &Expr) {
-    let globals = match &checker.semantic_model().scope().kind {
+    let scope = checker.semantic_model().scope();
-        ScopeKind::Class(class_def) => &class_def.globals,
+    if let Some(stmt) = scope.get_global(name) {
-        ScopeKind::Function(function_def) => &function_def.globals,
-        _ => return,
-    };
-    if let Some(stmt) = globals.get(name) {
         if expr.start() < stmt.start() {
             #[allow(deprecated)]
             let location = checker.locator.compute_source_location(stmt.start());

crates/ruff/src/rules/pyupgrade/rules/super_call_with_parameters.rs

@@ -2,7 +2,6 @@ use rustpython_parser::ast::{self, Arg, Expr, Ranged, Stmt};
 
 use ruff_diagnostics::{AlwaysAutofixableViolation, Diagnostic, Fix};
 use ruff_macros::{derive_message_formats, violation};
-use ruff_python_semantic::scope::ScopeKind;
 
 use crate::checkers::ast::Checker;
 use crate::registry::AsRule;
@@ -46,7 +45,7 @@ pub(crate) fn super_call_with_parameters(
     let scope = checker.semantic_model().scope();
 
     // Check: are we in a Function scope?
-    if !matches!(scope.kind, ScopeKind::Function(_)) {
+    if !scope.kind.is_any_function() {
         return;
     }
 

crates/ruff_python_ast/src/helpers.rs

@@ -961,18 +961,15 @@ where
 
 #[derive(Default)]
 struct GlobalStatementVisitor<'a> {
-    globals: FxHashMap<&'a str, &'a Stmt>,
+    globals: FxHashMap<&'a str, TextRange>,
 }
 
 impl<'a> StatementVisitor<'a> for GlobalStatementVisitor<'a> {
     fn visit_stmt(&mut self, stmt: &'a Stmt) {
         match stmt {
-            Stmt::Global(ast::StmtGlobal {
+            Stmt::Global(ast::StmtGlobal { names, range }) => {
-                names,
-                range: _range,
-            }) => {
                 for name in names {
-                    self.globals.insert(name.as_str(), stmt);
+                    self.globals.insert(name.as_str(), *range);
                 }
             }
             Stmt::FunctionDef(_) | Stmt::AsyncFunctionDef(_) | Stmt::ClassDef(_) => {
@@ -984,11 +981,9 @@ impl<'a> StatementVisitor<'a> for GlobalStatementVisitor<'a> {
 }
 
 /// Extract a map from global name to its last-defining [`Stmt`].
-pub fn extract_globals(body: &[Stmt]) -> FxHashMap<&str, &Stmt> {
+pub fn extract_globals(body: &[Stmt]) -> FxHashMap<&str, TextRange> {
     let mut visitor = GlobalStatementVisitor::default();
-    for stmt in body {
+    visitor.visit_body(body);
-        visitor.visit_stmt(stmt);
-    }
     visitor.globals
 }
 

crates/ruff_python_semantic/src/model.rs

@@ -212,7 +212,7 @@ impl<'a> SemanticModel<'a> {
                 }
             }
 
-            seen_function |= scope.kind.is_function();
+            seen_function |= scope.kind.is_any_function();
             import_starred = import_starred || scope.uses_star_imports();
         }
 
@@ -539,7 +539,7 @@ impl<'a> SemanticModel<'a> {
         self.scope_id == scope_id
     }
 
-    /// Return `true` if the context is at the top level of the module (i.e., in the module scope,
+    /// Return `true` if the model is at the top level of the module (i.e., in the module scope,
     /// and not nested within any statements).
     pub fn at_top_level(&self) -> bool {
         self.scope_id.is_global()
@@ -548,6 +548,19 @@ impl<'a> SemanticModel<'a> {
                 .map_or(true, |stmt_id| self.stmts.parent_id(stmt_id).is_none())
     }
 
+    /// Return `true` if the model is in an async context.
+    pub fn in_async_context(&self) -> bool {
+        for scope in self.scopes() {
+            if scope.kind.is_async_function() {
+                return true;
+            }
+            if scope.kind.is_function() {
+                return false;
+            }
+        }
+        false
+    }
+
     /// Returns `true` if the given [`BindingId`] is used.
     pub fn is_used(&self, binding_id: BindingId) -> bool {
         self.bindings[binding_id].is_used()

crates/ruff_python_semantic/src/scope.rs

@@ -1,8 +1,10 @@
 use std::ops::{Deref, DerefMut};
 
-use ruff_index::{newtype_index, Idx, IndexSlice, IndexVec};
+use ruff_text_size::TextRange;
 use rustc_hash::FxHashMap;
-use rustpython_parser::ast::{Arguments, Expr, Keyword, Stmt};
+use rustpython_parser::ast;
+
+use ruff_index::{newtype_index, Idx, IndexSlice, IndexVec};
 
 use crate::binding::{BindingId, StarImportation};
 
@@ -19,6 +21,8 @@ pub struct Scope<'a> {
     bindings: FxHashMap<&'a str, BindingId>,
     /// A map from bound name to binding index, for bindings that were shadowed later in the scope.
     shadowed_bindings: FxHashMap<&'a str, Vec<BindingId>>,
+    /// A map from global name to the range that declares it.
+    globals: FxHashMap<&'a str, TextRange>,
 }
 
 impl<'a> Scope<'a> {
@@ -30,6 +34,7 @@ impl<'a> Scope<'a> {
             star_imports: Vec::default(),
             bindings: FxHashMap::default(),
             shadowed_bindings: FxHashMap::default(),
+            globals: FxHashMap::default(),
         }
     }
 
@@ -41,6 +46,7 @@ impl<'a> Scope<'a> {
             star_imports: Vec::default(),
             bindings: FxHashMap::default(),
             shadowed_bindings: FxHashMap::default(),
+            globals: FxHashMap::default(),
         }
     }
 
@@ -103,48 +109,32 @@ impl<'a> Scope<'a> {
     pub fn star_imports(&self) -> impl Iterator<Item = &StarImportation<'a>> {
         self.star_imports.iter()
     }
+
+    /// Add a global name to this scope.
+    pub fn add_global(&mut self, name: &'a str, range: TextRange) {
+        self.globals.insert(name, range);
+    }
+
+    /// Returns the range of the global name with the given name.
+    pub fn get_global(&self, name: &str) -> Option<TextRange> {
+        self.globals.get(name).copied()
+    }
 }
 
 #[derive(Debug, is_macro::Is)]
 pub enum ScopeKind<'a> {
-    Class(ClassDef<'a>),
+    Class(&'a ast::StmtClassDef),
-    Function(FunctionDef<'a>),
+    Function(&'a ast::StmtFunctionDef),
+    AsyncFunction(&'a ast::StmtAsyncFunctionDef),
     Generator,
     Module,
-    Lambda(Lambda<'a>),
+    Lambda(&'a ast::ExprLambda),
 }
 
-#[derive(Debug)]
+impl ScopeKind<'_> {
-pub struct FunctionDef<'a> {
+    pub const fn is_any_function(&self) -> bool {
-    // Properties derived from Stmt::FunctionDef.
+        matches!(self, ScopeKind::Function(_) | ScopeKind::AsyncFunction(_))
-    pub name: &'a str,
+    }
-    pub args: &'a Arguments,
-    pub body: &'a [Stmt],
-    pub decorator_list: &'a [Expr],
-    // pub returns: Option<&'a Expr>,
-    // pub type_comment: Option<&'a str>,
-    // Scope-specific properties.
-    // TODO(charlie): Create AsyncFunctionDef to mirror the AST.
-    pub async_: bool,
-    pub globals: FxHashMap<&'a str, &'a Stmt>,
-}
-
-#[derive(Debug)]
-pub struct ClassDef<'a> {
-    // Properties derived from Stmt::ClassDef.
-    pub name: &'a str,
-    pub bases: &'a [Expr],
-    pub keywords: &'a [Keyword],
-    // pub body: &'a [Stmt],
-    pub decorator_list: &'a [Expr],
-    // Scope-specific properties.
-    pub globals: FxHashMap<&'a str, &'a Stmt>,
-}
-
-#[derive(Debug)]
-pub struct Lambda<'a> {
-    pub args: &'a Arguments,
-    pub body: &'a Expr,
 }
 
 /// Id uniquely identifying a scope in a program.