ty_python_core/
scope.rs

1use std::ops::Range;
2
3use ruff_db::{PythonFile, files::File, parsed::ParsedModuleRef};
4use ruff_index::newtype_index;
5use ruff_python_ast::{self as ast, NodeIndex};
6
7use crate::{
8    Db, Program, ProgramFile, SemanticIndex, ast_node_ref::AstNodeRef, definition::Definition,
9    node_key::NodeKey, semantic_index,
10};
11
12/// A cross-module identifier of a scope that can be used as a salsa query parameter.
13#[salsa::tracked(debug, heap_size=ruff_memory_usage::heap_size)]
14pub struct ScopeId<'db> {
15    #[returns(copy)]
16    pub program_file: ProgramFile<'db>,
17
18    #[returns(copy)]
19    pub file_scope_id: FileScopeId,
20}
21
22// The Salsa heap is tracked separately.
23impl get_size2::GetSize for ScopeId<'_> {}
24
25impl<'db> ScopeId<'db> {
26    pub fn file(self, db: &dyn Db) -> File {
27        self.program_file(db).file(db)
28    }
29
30    pub fn python_file(self, db: &'db dyn Db) -> PythonFile<'db> {
31        self.program_file(db).python_file(db)
32    }
33
34    pub fn program(self, db: &'db dyn Db) -> Program<'db> {
35        self.program_file(db).program(db)
36    }
37
38    pub fn is_annotation(self, db: &'db dyn Db) -> bool {
39        self.node(db).scope_kind().is_annotation()
40    }
41
42    pub fn node(self, db: &'db dyn Db) -> &'db NodeWithScopeKind {
43        self.scope(db).node()
44    }
45
46    /// Returns `true` if this scope may require type context from its parent scope.
47    pub fn accepts_type_context(self, db: &'db dyn Db) -> bool {
48        matches!(
49            self.node(db),
50            NodeWithScopeKind::Lambda(_)
51                | NodeWithScopeKind::ListComprehension(_)
52                | NodeWithScopeKind::SetComprehension(_)
53                | NodeWithScopeKind::DictComprehension(_)
54                | NodeWithScopeKind::GeneratorExpression(_)
55        )
56    }
57
58    pub fn scope(self, db: &'db dyn Db) -> &'db Scope {
59        semantic_index(db, self.program_file(db)).scope(self.file_scope_id(db))
60    }
61
62    /// Returns the class definition for the enclosing class if this scope is a method body.
63    pub fn class_definition_of_method(self, db: &'db dyn Db) -> Option<Definition<'db>> {
64        semantic_index(db, self.program_file(db)).class_definition_of_method(self.file_scope_id(db))
65    }
66
67    pub fn is_method_scope(self, db: &'db dyn Db) -> bool {
68        self.class_definition_of_method(db).is_some()
69    }
70
71    pub fn name<'ast>(self, db: &'db dyn Db, module: &'ast ParsedModuleRef) -> &'ast str {
72        match self.node(db) {
73            NodeWithScopeKind::Module => "<module>",
74            NodeWithScopeKind::Class(class) | NodeWithScopeKind::ClassTypeParameters(class) => {
75                class.node(module).name.as_str()
76            }
77            NodeWithScopeKind::Function(function)
78            | NodeWithScopeKind::FunctionTypeParameters(function) => {
79                function.node(module).name.as_str()
80            }
81            NodeWithScopeKind::TypeAlias(type_alias)
82            | NodeWithScopeKind::TypeAliasTypeParameters(type_alias) => type_alias
83                .node(module)
84                .name
85                .as_name_expr()
86                .map(|name| name.id.as_str())
87                .unwrap_or("<type alias>"),
88            NodeWithScopeKind::Lambda(_) => "<lambda>",
89            NodeWithScopeKind::ListComprehension(_) => "<listcomp>",
90            NodeWithScopeKind::SetComprehension(_) => "<setcomp>",
91            NodeWithScopeKind::DictComprehension(_) => "<dictcomp>",
92            NodeWithScopeKind::GeneratorExpression(_) => "<generator>",
93        }
94    }
95}
96
97/// ID that uniquely identifies a scope inside of a module.
98#[newtype_index]
99#[derive(Ord, PartialOrd, get_size2::GetSize)]
100pub struct FileScopeId;
101
102impl FileScopeId {
103    /// Returns the scope id of the module-global scope.
104    pub fn global() -> Self {
105        FileScopeId::from_u32(0)
106    }
107
108    pub fn is_global(self) -> bool {
109        self == FileScopeId::global()
110    }
111
112    pub fn to_scope_id<'db>(self, db: &'db dyn Db, file: ProgramFile<'db>) -> ScopeId<'db> {
113        let index = semantic_index(db, file);
114        index.scope_ids_by_scope[self]
115    }
116
117    pub fn is_generator_function(self, index: &SemanticIndex) -> bool {
118        index.generator_functions.contains(&self)
119    }
120
121    pub fn is_async_comprehension(self, index: &SemanticIndex) -> bool {
122        index.async_comprehensions.contains(&self)
123    }
124}
125
126#[derive(Debug, get_size2::GetSize)]
127pub struct Scope {
128    /// The parent scope, if any.
129    parent: Option<FileScopeId>,
130
131    /// The node that introduces this scope.
132    node: NodeWithScopeKind,
133
134    /// The range of [`FileScopeId`]s that are descendants of this scope.
135    descendants: Range<FileScopeId>,
136}
137
138impl Scope {
139    pub(super) fn new(
140        parent: Option<FileScopeId>,
141        node: NodeWithScopeKind,
142        descendants: Range<FileScopeId>,
143    ) -> Self {
144        Scope {
145            parent,
146            node,
147            descendants,
148        }
149    }
150
151    pub fn parent(&self) -> Option<FileScopeId> {
152        self.parent
153    }
154
155    pub fn node(&self) -> &NodeWithScopeKind {
156        &self.node
157    }
158
159    pub fn kind(&self) -> ScopeKind {
160        self.node().scope_kind()
161    }
162
163    pub fn visibility(&self) -> ScopeVisibility {
164        self.kind().visibility()
165    }
166
167    pub(crate) fn descendants(&self) -> Range<FileScopeId> {
168        self.descendants.clone()
169    }
170
171    pub(super) fn extend_descendants(&mut self, children_end: FileScopeId) {
172        self.descendants = self.descendants.start..children_end;
173    }
174
175    pub fn is_eager(&self) -> bool {
176        self.kind().is_eager()
177    }
178}
179
180#[derive(Debug, PartialEq, Eq, Clone, Copy, Hash, get_size2::GetSize)]
181pub enum ScopeVisibility {
182    /// The scope is private (e.g. function, type alias, comprehension scope).
183    Private,
184    /// The scope is public (e.g. module, class scope).
185    Public,
186}
187
188impl ScopeVisibility {
189    pub(crate) const fn is_public(self) -> bool {
190        matches!(self, ScopeVisibility::Public)
191    }
192
193    pub const fn is_private(self) -> bool {
194        matches!(self, ScopeVisibility::Private)
195    }
196}
197
198#[derive(Debug, PartialEq, Eq, PartialOrd, Ord, Clone, Copy, Hash, get_size2::GetSize)]
199pub(crate) enum ScopeLaziness {
200    /// The scope is evaluated lazily (e.g. function, type alias scope).
201    Lazy,
202    /// The scope is evaluated eagerly (e.g. module, class, comprehension scope).
203    Eager,
204}
205
206impl ScopeLaziness {
207    pub(crate) const fn is_eager(self) -> bool {
208        matches!(self, ScopeLaziness::Eager)
209    }
210
211    pub(crate) const fn is_lazy(self) -> bool {
212        matches!(self, ScopeLaziness::Lazy)
213    }
214}
215
216#[derive(Copy, Clone, Debug, PartialEq, Eq)]
217pub enum ScopeKind {
218    Module,
219    TypeParams,
220    Class,
221    Function,
222    Lambda,
223    Comprehension,
224    TypeAlias,
225}
226
227impl ScopeKind {
228    pub(crate) const fn is_eager(self) -> bool {
229        self.laziness().is_eager()
230    }
231
232    pub(crate) const fn laziness(self) -> ScopeLaziness {
233        match self {
234            ScopeKind::Module
235            | ScopeKind::Class
236            | ScopeKind::Comprehension
237            | ScopeKind::TypeParams => ScopeLaziness::Eager,
238            ScopeKind::Function | ScopeKind::Lambda | ScopeKind::TypeAlias => ScopeLaziness::Lazy,
239        }
240    }
241
242    const fn visibility(self) -> ScopeVisibility {
243        match self {
244            ScopeKind::Module | ScopeKind::Class => ScopeVisibility::Public,
245            ScopeKind::TypeParams
246            | ScopeKind::TypeAlias
247            | ScopeKind::Function
248            | ScopeKind::Lambda
249            | ScopeKind::Comprehension => ScopeVisibility::Private,
250        }
251    }
252
253    pub const fn is_function_like(self) -> bool {
254        // Type parameter scopes behave like function scopes in terms of name resolution; CPython
255        // symbol table also uses the term "function-like" for these scopes.
256        matches!(
257            self,
258            ScopeKind::TypeParams
259                | ScopeKind::Function
260                | ScopeKind::Lambda
261                | ScopeKind::TypeAlias
262                | ScopeKind::Comprehension
263        )
264    }
265
266    pub const fn is_class(self) -> bool {
267        matches!(self, ScopeKind::Class)
268    }
269
270    pub const fn is_module(self) -> bool {
271        matches!(self, ScopeKind::Module)
272    }
273
274    pub(crate) const fn is_annotation(self) -> bool {
275        matches!(self, ScopeKind::TypeParams | ScopeKind::TypeAlias)
276    }
277
278    pub const fn is_non_lambda_function(self) -> bool {
279        matches!(self, ScopeKind::Function)
280    }
281}
282
283/// Reference to a node that introduces a new scope.
284#[derive(Copy, Clone, Debug)]
285pub enum NodeWithScopeRef<'a> {
286    Module,
287    Class(&'a ast::StmtClassDef),
288    Function(&'a ast::StmtFunctionDef),
289    Lambda(&'a ast::ExprLambda),
290    FunctionTypeParameters(&'a ast::StmtFunctionDef),
291    ClassTypeParameters(&'a ast::StmtClassDef),
292    TypeAlias(&'a ast::StmtTypeAlias),
293    TypeAliasTypeParameters(&'a ast::StmtTypeAlias),
294    ListComprehension(&'a ast::ExprListComp),
295    SetComprehension(&'a ast::ExprSetComp),
296    DictComprehension(&'a ast::ExprDictComp),
297    GeneratorExpression(&'a ast::ExprGenerator),
298}
299
300impl NodeWithScopeRef<'_> {
301    /// Converts the unowned reference to an owned [`NodeWithScopeKind`].
302    ///
303    /// Note that node wrapped by `self` must be a child of `module`.
304    pub(super) fn to_kind(self, module: &ParsedModuleRef) -> NodeWithScopeKind {
305        match self {
306            NodeWithScopeRef::Module => NodeWithScopeKind::Module,
307            NodeWithScopeRef::Class(class) => {
308                NodeWithScopeKind::Class(AstNodeRef::new(module, class))
309            }
310            NodeWithScopeRef::Function(function) => {
311                NodeWithScopeKind::Function(AstNodeRef::new(module, function))
312            }
313            NodeWithScopeRef::TypeAlias(type_alias) => {
314                NodeWithScopeKind::TypeAlias(AstNodeRef::new(module, type_alias))
315            }
316            NodeWithScopeRef::TypeAliasTypeParameters(type_alias) => {
317                NodeWithScopeKind::TypeAliasTypeParameters(AstNodeRef::new(module, type_alias))
318            }
319            NodeWithScopeRef::Lambda(lambda) => {
320                NodeWithScopeKind::Lambda(AstNodeRef::new(module, lambda))
321            }
322            NodeWithScopeRef::FunctionTypeParameters(function) => {
323                NodeWithScopeKind::FunctionTypeParameters(AstNodeRef::new(module, function))
324            }
325            NodeWithScopeRef::ClassTypeParameters(class) => {
326                NodeWithScopeKind::ClassTypeParameters(AstNodeRef::new(module, class))
327            }
328            NodeWithScopeRef::ListComprehension(comprehension) => {
329                NodeWithScopeKind::ListComprehension(AstNodeRef::new(module, comprehension))
330            }
331            NodeWithScopeRef::SetComprehension(comprehension) => {
332                NodeWithScopeKind::SetComprehension(AstNodeRef::new(module, comprehension))
333            }
334            NodeWithScopeRef::DictComprehension(comprehension) => {
335                NodeWithScopeKind::DictComprehension(AstNodeRef::new(module, comprehension))
336            }
337            NodeWithScopeRef::GeneratorExpression(generator) => {
338                NodeWithScopeKind::GeneratorExpression(AstNodeRef::new(module, generator))
339            }
340        }
341    }
342
343    pub(crate) fn node_key(self) -> NodeWithScopeKey {
344        match self {
345            NodeWithScopeRef::Module => NodeWithScopeKey::Module,
346            NodeWithScopeRef::Class(class) => NodeWithScopeKey::Class(NodeKey::from_node(class)),
347            NodeWithScopeRef::Function(function) => {
348                NodeWithScopeKey::Function(NodeKey::from_node(function))
349            }
350            NodeWithScopeRef::Lambda(lambda) => {
351                NodeWithScopeKey::Lambda(NodeKey::from_node(lambda))
352            }
353            NodeWithScopeRef::FunctionTypeParameters(function) => {
354                NodeWithScopeKey::FunctionTypeParameters(NodeKey::from_node(function))
355            }
356            NodeWithScopeRef::ClassTypeParameters(class) => {
357                NodeWithScopeKey::ClassTypeParameters(NodeKey::from_node(class))
358            }
359            NodeWithScopeRef::TypeAlias(type_alias) => {
360                NodeWithScopeKey::TypeAlias(NodeKey::from_node(type_alias))
361            }
362            NodeWithScopeRef::TypeAliasTypeParameters(type_alias) => {
363                NodeWithScopeKey::TypeAliasTypeParameters(NodeKey::from_node(type_alias))
364            }
365            NodeWithScopeRef::ListComprehension(comprehension) => {
366                NodeWithScopeKey::ListComprehension(NodeKey::from_node(comprehension))
367            }
368            NodeWithScopeRef::SetComprehension(comprehension) => {
369                NodeWithScopeKey::SetComprehension(NodeKey::from_node(comprehension))
370            }
371            NodeWithScopeRef::DictComprehension(comprehension) => {
372                NodeWithScopeKey::DictComprehension(NodeKey::from_node(comprehension))
373            }
374            NodeWithScopeRef::GeneratorExpression(generator) => {
375                NodeWithScopeKey::GeneratorExpression(NodeKey::from_node(generator))
376            }
377        }
378    }
379}
380
381/// Node that introduces a new scope.
382#[derive(Clone, Debug, get_size2::GetSize)]
383pub enum NodeWithScopeKind {
384    Module,
385    Class(AstNodeRef<ast::StmtClassDef>),
386    ClassTypeParameters(AstNodeRef<ast::StmtClassDef>),
387    Function(AstNodeRef<ast::StmtFunctionDef>),
388    FunctionTypeParameters(AstNodeRef<ast::StmtFunctionDef>),
389    TypeAliasTypeParameters(AstNodeRef<ast::StmtTypeAlias>),
390    TypeAlias(AstNodeRef<ast::StmtTypeAlias>),
391    Lambda(AstNodeRef<ast::ExprLambda>),
392    ListComprehension(AstNodeRef<ast::ExprListComp>),
393    SetComprehension(AstNodeRef<ast::ExprSetComp>),
394    DictComprehension(AstNodeRef<ast::ExprDictComp>),
395    GeneratorExpression(AstNodeRef<ast::ExprGenerator>),
396}
397
398impl NodeWithScopeKind {
399    pub const fn scope_kind(&self) -> ScopeKind {
400        match self {
401            Self::Module => ScopeKind::Module,
402            Self::Class(_) => ScopeKind::Class,
403            Self::Function(_) => ScopeKind::Function,
404            Self::Lambda(_) => ScopeKind::Lambda,
405            Self::FunctionTypeParameters(_)
406            | Self::ClassTypeParameters(_)
407            | Self::TypeAliasTypeParameters(_) => ScopeKind::TypeParams,
408            Self::TypeAlias(_) => ScopeKind::TypeAlias,
409            Self::ListComprehension(_)
410            | Self::SetComprehension(_)
411            | Self::DictComprehension(_)
412            | Self::GeneratorExpression(_) => ScopeKind::Comprehension,
413        }
414    }
415
416    pub fn as_class(&self) -> Option<&AstNodeRef<ast::StmtClassDef>> {
417        match self {
418            Self::Class(class) => Some(class),
419            _ => None,
420        }
421    }
422
423    pub fn expect_class(&self) -> &AstNodeRef<ast::StmtClassDef> {
424        self.as_class().expect("expected class")
425    }
426
427    pub fn as_function(&self) -> Option<&AstNodeRef<ast::StmtFunctionDef>> {
428        match self {
429            Self::Function(function) => Some(function),
430            _ => None,
431        }
432    }
433
434    pub fn expect_function(&self) -> &AstNodeRef<ast::StmtFunctionDef> {
435        self.as_function().expect("expected function")
436    }
437
438    fn as_type_alias(&self) -> Option<&AstNodeRef<ast::StmtTypeAlias>> {
439        match self {
440            Self::TypeAlias(type_alias) => Some(type_alias),
441            _ => None,
442        }
443    }
444
445    pub fn expect_type_alias(&self) -> &AstNodeRef<ast::StmtTypeAlias> {
446        self.as_type_alias().expect("expected type alias")
447    }
448
449    /// Returns the anchor node index for this scope, or `None` for the module scope.
450    ///
451    /// This is used to compute relative node indices for expressions within the scope,
452    /// providing a stable anchor that only changes when the scope-introducing node changes.
453    pub fn node_index(&self) -> Option<NodeIndex> {
454        match self {
455            Self::Module => None,
456            Self::Class(class) => Some(class.index()),
457            Self::ClassTypeParameters(class) => Some(class.index()),
458            Self::Function(function) => Some(function.index()),
459            Self::FunctionTypeParameters(function) => Some(function.index()),
460            Self::TypeAlias(type_alias) => Some(type_alias.index()),
461            Self::TypeAliasTypeParameters(type_alias) => Some(type_alias.index()),
462            Self::Lambda(lambda) => Some(lambda.index()),
463            Self::ListComprehension(comp) => Some(comp.index()),
464            Self::SetComprehension(comp) => Some(comp.index()),
465            Self::DictComprehension(comp) => Some(comp.index()),
466            Self::GeneratorExpression(generator) => Some(generator.index()),
467        }
468    }
469}
470
471#[derive(Copy, Clone, Debug, Eq, PartialEq, Hash, get_size2::GetSize)]
472pub enum NodeWithScopeKey {
473    Module,
474    Class(NodeKey),
475    ClassTypeParameters(NodeKey),
476    Function(NodeKey),
477    FunctionTypeParameters(NodeKey),
478    TypeAlias(NodeKey),
479    TypeAliasTypeParameters(NodeKey),
480    Lambda(NodeKey),
481    ListComprehension(NodeKey),
482    SetComprehension(NodeKey),
483    DictComprehension(NodeKey),
484    GeneratorExpression(NodeKey),
485}
ty_python_core/scope.rs

ty_python_core/
scope.rs
