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

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>),
-    Function(FunctionDef<'a>),
+    Class(&'a ast::StmtClassDef),
+    Function(&'a ast::StmtFunctionDef),
+    AsyncFunction(&'a ast::StmtAsyncFunctionDef),
     Generator,
     Module,
-    Lambda(Lambda<'a>),
+    Lambda(&'a ast::ExprLambda),
 }
 
-#[derive(Debug)]
-pub struct FunctionDef<'a> {
-    // Properties derived from Stmt::FunctionDef.
-    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,
+impl ScopeKind<'_> {
+    pub const fn is_any_function(&self) -> bool {
+        matches!(self, ScopeKind::Function(_) | ScopeKind::AsyncFunction(_))
+    }
 }
 
 /// Id uniquely identifying a scope in a program.