[ty] add support for nonlocal statements

This commit is contained in:
Jack O'Connor 2025-07-02 19:08:32 -07:00
parent 110765154f
commit 78bd73f25a
11 changed files with 625 additions and 44 deletions

View file

@ -1421,7 +1421,7 @@ impl RequiresExplicitReExport {
/// ```py
/// def _():
/// x = 1
///
///
/// x = 2
///
/// if flag():

View file

@ -217,9 +217,6 @@ pub(crate) struct SemanticIndex<'db> {
/// Map from the file-local [`FileScopeId`] to the salsa-ingredient [`ScopeId`].
scope_ids_by_scope: IndexVec<FileScopeId, ScopeId<'db>>,
/// Map from the file-local [`FileScopeId`] to the set of explicit-global symbols it contains.
globals_by_scope: FxHashMap<FileScopeId, FxHashSet<ScopedPlaceId>>,
/// Use-def map for each scope in this file.
use_def_maps: IndexVec<FileScopeId, ArcUseDefMap<'db>>,
@ -308,9 +305,19 @@ impl<'db> SemanticIndex<'db> {
symbol: ScopedPlaceId,
scope: FileScopeId,
) -> bool {
self.globals_by_scope
.get(&scope)
.is_some_and(|globals| globals.contains(&symbol))
self.place_table(scope)
.place_expr(symbol)
.is_marked_global()
}
pub(crate) fn symbol_is_nonlocal_in_scope(
&self,
symbol: ScopedPlaceId,
scope: FileScopeId,
) -> bool {
self.place_table(scope)
.place_expr(symbol)
.is_marked_nonlocal()
}
/// Returns the id of the parent scope.

View file

@ -103,7 +103,6 @@ pub(super) struct SemanticIndexBuilder<'db, 'ast> {
use_def_maps: IndexVec<FileScopeId, UseDefMapBuilder<'db>>,
scopes_by_node: FxHashMap<NodeWithScopeKey, FileScopeId>,
scopes_by_expression: FxHashMap<ExpressionNodeKey, FileScopeId>,
globals_by_scope: FxHashMap<FileScopeId, FxHashSet<ScopedPlaceId>>,
definitions_by_node: FxHashMap<DefinitionNodeKey, Definitions<'db>>,
expressions_by_node: FxHashMap<ExpressionNodeKey, Expression<'db>>,
imported_modules: FxHashSet<ModuleName>,
@ -141,7 +140,6 @@ impl<'db, 'ast> SemanticIndexBuilder<'db, 'ast> {
scopes_by_node: FxHashMap::default(),
definitions_by_node: FxHashMap::default(),
expressions_by_node: FxHashMap::default(),
globals_by_scope: FxHashMap::default(),
imported_modules: FxHashSet::default(),
generator_functions: FxHashSet::default(),
@ -349,7 +347,12 @@ impl<'db, 'ast> SemanticIndexBuilder<'db, 'ast> {
popped_scope_id
}
fn current_place_table(&mut self) -> &mut PlaceTableBuilder {
fn current_place_table(&self) -> &PlaceTableBuilder {
let scope_id = self.current_scope();
&self.place_tables[scope_id]
}
fn current_place_table_mut(&mut self) -> &mut PlaceTableBuilder {
let scope_id = self.current_scope();
&mut self.place_tables[scope_id]
}
@ -389,7 +392,7 @@ impl<'db, 'ast> SemanticIndexBuilder<'db, 'ast> {
/// Add a symbol to the place table and the use-def map.
/// Return the [`ScopedPlaceId`] that uniquely identifies the symbol in both.
fn add_symbol(&mut self, name: Name) -> ScopedPlaceId {
let (place_id, added) = self.current_place_table().add_symbol(name);
let (place_id, added) = self.current_place_table_mut().add_symbol(name);
if added {
self.current_use_def_map_mut().add_place(place_id);
}
@ -399,7 +402,7 @@ impl<'db, 'ast> SemanticIndexBuilder<'db, 'ast> {
/// Add a place to the place table and the use-def map.
/// Return the [`ScopedPlaceId`] that uniquely identifies the place in both.
fn add_place(&mut self, place_expr: PlaceExprWithFlags) -> ScopedPlaceId {
let (place_id, added) = self.current_place_table().add_place(place_expr);
let (place_id, added) = self.current_place_table_mut().add_place(place_expr);
if added {
self.current_use_def_map_mut().add_place(place_id);
}
@ -407,15 +410,15 @@ impl<'db, 'ast> SemanticIndexBuilder<'db, 'ast> {
}
fn mark_place_bound(&mut self, id: ScopedPlaceId) {
self.current_place_table().mark_place_bound(id);
self.current_place_table_mut().mark_place_bound(id);
}
fn mark_place_declared(&mut self, id: ScopedPlaceId) {
self.current_place_table().mark_place_declared(id);
self.current_place_table_mut().mark_place_declared(id);
}
fn mark_place_used(&mut self, id: ScopedPlaceId) {
self.current_place_table().mark_place_used(id);
self.current_place_table_mut().mark_place_used(id);
}
fn add_entry_for_definition_key(&mut self, key: DefinitionNodeKey) -> &mut Definitions<'db> {
@ -1042,7 +1045,6 @@ impl<'db, 'ast> SemanticIndexBuilder<'db, 'ast> {
self.scopes_by_node.shrink_to_fit();
self.generator_functions.shrink_to_fit();
self.eager_snapshots.shrink_to_fit();
self.globals_by_scope.shrink_to_fit();
SemanticIndex {
place_tables,
@ -1050,7 +1052,6 @@ impl<'db, 'ast> SemanticIndexBuilder<'db, 'ast> {
definitions_by_node: self.definitions_by_node,
expressions_by_node: self.expressions_by_node,
scope_ids_by_scope: self.scope_ids_by_scope,
globals_by_scope: self.globals_by_scope,
ast_ids,
scopes_by_expression: self.scopes_by_expression,
scopes_by_node: self.scopes_by_node,
@ -1418,6 +1419,29 @@ impl<'ast> Visitor<'ast> for SemanticIndexBuilder<'_, 'ast> {
self.visit_expr(value);
}
if let ast::Expr::Name(name) = &*node.target {
let symbol_id = self.add_symbol(name.id.clone());
let symbol = self.current_place_table().place_expr(symbol_id);
// Check whether the variable has been declared global.
if symbol.is_marked_global() {
self.report_semantic_error(SemanticSyntaxError {
kind: SemanticSyntaxErrorKind::AnnotatedGlobal(name.id.as_str().into()),
range: name.range,
python_version: self.python_version,
});
}
// Check whether the variable has been declared nonlocal.
if symbol.is_marked_nonlocal() {
self.report_semantic_error(SemanticSyntaxError {
kind: SemanticSyntaxErrorKind::AnnotatedNonlocal(
name.id.as_str().into(),
),
range: name.range,
python_version: self.python_version,
});
}
}
// See https://docs.python.org/3/library/ast.html#ast.AnnAssign
if matches!(
*node.target,
@ -1858,8 +1882,8 @@ impl<'ast> Visitor<'ast> for SemanticIndexBuilder<'_, 'ast> {
}) => {
for name in names {
let symbol_id = self.add_symbol(name.id.clone());
let symbol_table = self.current_place_table();
let symbol = symbol_table.place_expr(symbol_id);
let symbol = self.current_place_table().place_expr(symbol_id);
// Check whether the variable has already been accessed in this scope.
if symbol.is_bound() || symbol.is_declared() || symbol.is_used() {
self.report_semantic_error(SemanticSyntaxError {
kind: SemanticSyntaxErrorKind::LoadBeforeGlobalDeclaration {
@ -1870,11 +1894,56 @@ impl<'ast> Visitor<'ast> for SemanticIndexBuilder<'_, 'ast> {
python_version: self.python_version,
});
}
let scope_id = self.current_scope();
self.globals_by_scope
.entry(scope_id)
.or_default()
.insert(symbol_id);
// Check whether the variable has also been declared nonlocal.
if symbol.is_marked_nonlocal() {
self.report_semantic_error(SemanticSyntaxError {
kind: SemanticSyntaxErrorKind::NonlocalAndGlobal(name.to_string()),
range: name.range,
python_version: self.python_version,
});
}
self.current_place_table_mut().mark_place_global(symbol_id);
}
walk_stmt(self, stmt);
}
ast::Stmt::Nonlocal(ast::StmtNonlocal {
range: _,
node_index: _,
names,
}) => {
for name in names {
let symbol_id = self.add_symbol(name.id.clone());
let symbol = self.current_place_table().place_expr(symbol_id);
// Check whether the variable has already been accessed in this scope.
if symbol.is_bound() || symbol.is_declared() || symbol.is_used() {
self.report_semantic_error(SemanticSyntaxError {
kind: SemanticSyntaxErrorKind::LoadBeforeNonlocalDeclaration {
name: name.to_string(),
start: name.range.start(),
},
range: name.range,
python_version: self.python_version,
});
}
// Check whether the variable has also been declared global.
if symbol.is_marked_global() {
self.report_semantic_error(SemanticSyntaxError {
kind: SemanticSyntaxErrorKind::NonlocalAndGlobal(name.to_string()),
range: name.range,
python_version: self.python_version,
});
}
// The variable is required to exist in an enclosing scope, but that definition
// might come later. For example, this is example legal, but we can't check
// that here, because we haven't gotten to `x = 1`:
// ```py
// def f():
// def g():
// nonlocal x
// x = 1
// ```
self.current_place_table_mut()
.mark_place_nonlocal(symbol_id);
}
walk_stmt(self, stmt);
}
@ -1888,7 +1957,7 @@ impl<'ast> Visitor<'ast> for SemanticIndexBuilder<'_, 'ast> {
for target in targets {
if let Ok(target) = PlaceExpr::try_from(target) {
let place_id = self.add_place(PlaceExprWithFlags::new(target));
self.current_place_table().mark_place_used(place_id);
self.current_place_table_mut().mark_place_used(place_id);
self.delete_binding(place_id);
}
}

View file

@ -330,6 +330,16 @@ impl PlaceExprWithFlags {
self.flags.contains(PlaceFlags::IS_DECLARED)
}
/// Is the place `global` its containing scope?
pub fn is_marked_global(&self) -> bool {
self.flags.contains(PlaceFlags::MARKED_GLOBAL)
}
/// Is the place `nonlocal` its containing scope?
pub fn is_marked_nonlocal(&self) -> bool {
self.flags.contains(PlaceFlags::MARKED_NONLOCAL)
}
pub(crate) fn as_name(&self) -> Option<&Name> {
self.expr.as_name()
}
@ -397,9 +407,7 @@ bitflags! {
const IS_USED = 1 << 0;
const IS_BOUND = 1 << 1;
const IS_DECLARED = 1 << 2;
/// TODO: This flag is not yet set by anything
const MARKED_GLOBAL = 1 << 3;
/// TODO: This flag is not yet set by anything
const MARKED_NONLOCAL = 1 << 4;
const IS_INSTANCE_ATTRIBUTE = 1 << 5;
}
@ -663,7 +671,7 @@ impl PlaceTable {
}
/// Returns the place named `name`.
#[allow(unused)] // used in tests
#[cfg(test)]
pub(crate) fn place_by_name(&self, name: &str) -> Option<&PlaceExprWithFlags> {
let id = self.place_id_by_name(name)?;
Some(self.place_expr(id))
@ -814,6 +822,14 @@ impl PlaceTableBuilder {
self.table.places[id].insert_flags(PlaceFlags::IS_USED);
}
pub(super) fn mark_place_global(&mut self, id: ScopedPlaceId) {
self.table.places[id].insert_flags(PlaceFlags::MARKED_GLOBAL);
}
pub(super) fn mark_place_nonlocal(&mut self, id: ScopedPlaceId) {
self.table.places[id].insert_flags(PlaceFlags::MARKED_NONLOCAL);
}
pub(super) fn places(&self) -> impl Iterator<Item = &PlaceExprWithFlags> {
self.table.places()
}

View file

@ -35,6 +35,7 @@
//! be considered a bug.)
use itertools::{Either, Itertools};
use ruff_db::diagnostic::{Annotation, DiagnosticId, Severity};
use ruff_db::files::File;
use ruff_db::parsed::{ParsedModuleRef, parsed_module};
use ruff_python_ast::visitor::{Visitor, walk_expr};
@ -1562,6 +1563,7 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> {
let mut bound_ty = ty;
let global_use_def_map = self.index.use_def_map(FileScopeId::global());
let nonlocal_use_def_map;
let place_id = binding.place(self.db());
let place = place_table.place_expr(place_id);
let skip_non_global_scopes = self.skip_non_global_scopes(file_scope_id, place_id);
@ -1572,9 +1574,58 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> {
.place_id_by_expr(&place.expr)
{
Some(id) => (global_use_def_map.end_of_scope_declarations(id), false),
// This case is a syntax error (load before global declaration) but ignore that here
// This variable shows up in `global` declarations but doesn't have an explicit
// binding in the global scope.
None => (use_def.declarations_at_binding(binding), true),
}
} else if self
.index
.symbol_is_nonlocal_in_scope(place_id, file_scope_id)
{
// If we run out of ancestor scopes without finding a definition, we'll fall back to
// the local scope. This will also be a syntax error in `infer_nonlocal_statement` (no
// binding for `nonlocal` found), but ignore that here.
let mut declarations = use_def.declarations_at_binding(binding);
let mut is_local = true;
// Walk up parent scopes looking for the enclosing scope that has definition of this
// name. `ancestor_scopes` includes the current scope, so skip that one.
for (enclosing_scope_file_id, enclosing_scope) in
self.index.ancestor_scopes(file_scope_id).skip(1)
{
// Ignore class scopes and the global scope.
if !enclosing_scope.kind().is_function_like() {
continue;
}
let enclosing_place_table = self.index.place_table(enclosing_scope_file_id);
let Some(enclosing_place_id) = enclosing_place_table.place_id_by_expr(&place.expr)
else {
// This ancestor scope doesn't have a binding. Keep going.
continue;
};
if self
.index
.symbol_is_nonlocal_in_scope(enclosing_place_id, enclosing_scope_file_id)
{
// The variable is `nonlocal` in this ancestor scope. Keep going.
continue;
}
if self
.index
.symbol_is_global_in_scope(enclosing_place_id, enclosing_scope_file_id)
{
// The variable is `global` in this ancestor scope. This breaks the `nonlocal`
// chain, and it's a syntax error in `infer_nonlocal_statement`. Ignore that
// here and just bail out of this loop.
break;
}
// We found the closest definition. Note that (unlike in `infer_place_load`) this
// does *not* need to be a binding. It could be just `x: int`.
nonlocal_use_def_map = self.index.use_def_map(enclosing_scope_file_id);
declarations = nonlocal_use_def_map.end_of_scope_declarations(enclosing_place_id);
is_local = false;
break;
}
(declarations, is_local)
} else {
(use_def.declarations_at_binding(binding), true)
};
@ -2204,12 +2255,12 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> {
ast::Stmt::Raise(raise) => self.infer_raise_statement(raise),
ast::Stmt::Return(ret) => self.infer_return_statement(ret),
ast::Stmt::Delete(delete) => self.infer_delete_statement(delete),
ast::Stmt::Nonlocal(nonlocal) => self.infer_nonlocal_statement(nonlocal),
ast::Stmt::Break(_)
| ast::Stmt::Continue(_)
| ast::Stmt::Pass(_)
| ast::Stmt::IpyEscapeCommand(_)
| ast::Stmt::Global(_)
| ast::Stmt::Nonlocal(_) => {
| ast::Stmt::Global(_) => {
// No-op
}
}
@ -4609,6 +4660,69 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> {
}
}
fn infer_nonlocal_statement(&mut self, nonlocal: &ast::StmtNonlocal) {
let ast::StmtNonlocal {
node_index: _,
range,
names,
} = nonlocal;
let db = self.db();
let scope = self.scope();
let file_scope_id = scope.file_scope_id(db);
let current_file = self.file();
'names: for name in names {
// Walk up parent scopes looking for a possible enclosing scope that may have a
// definition of this name visible to us. Note that we skip the scope containing the
// use that we are resolving, since we already looked for the place there up above.
for (enclosing_scope_file_id, _) in self.index.ancestor_scopes(file_scope_id).skip(1) {
// Class scopes are not visible to nested scopes, and `nonlocal` cannot refer to
// globals, so check only function-like scopes.
let enclosing_scope_id = enclosing_scope_file_id.to_scope_id(db, current_file);
if !enclosing_scope_id.is_function_like(db) {
continue;
}
let enclosing_place_table = self.index.place_table(enclosing_scope_file_id);
let Some(enclosing_place_id) = enclosing_place_table.place_id_by_name(name) else {
// This scope doesn't define this name. Keep going.
continue;
};
// We've found a definition for this name in an enclosing function-like scope.
// Either this definition is the valid place this name refers to, or else we'll
// emit a syntax error. Either way, we won't walk any more enclosing scopes. Note
// that there are differences here compared to `infer_place_load`: A regular load
// (e.g. `print(x)`) is allowed to refer to a global variable (e.g. `x = 1` in the
// global scope), and similarly it's allowed to refer to a local variable in an
// enclosing function that's declared `global` (e.g. `global x`). However, the
// `nonlocal` keyword can't refer to global variables (that's a `SyntaxError`), and
// it also can't refer to local variables in enclosing functions that are declared
// `global` (also a `SyntaxError`).
if self
.index
.symbol_is_global_in_scope(enclosing_place_id, enclosing_scope_file_id)
{
// A "chain" of `nonlocal` statements is "broken" by a `global` statement. Stop
// looping and report that this `nonlocal` statement is invalid.
break;
}
// We found a definition. We've checked that the name isn't `global` in this scope,
// but it's ok if it's `nonlocal`. If a "chain" of `nonlocal` statements fails to
// lead to a valid binding, the outermost one will be an error; we don't need to
// walk the whole chain for each one.
continue 'names;
}
// There's no matching binding in an enclosing scope. This `nonlocal` statement is
// invalid.
if let Some(builder) = self
.context
.report_diagnostic(DiagnosticId::InvalidSyntax, Severity::Error)
{
builder
.into_diagnostic(format_args!("no binding for nonlocal `{name}` found"))
.annotate(Annotation::primary(self.context.span(*range)));
}
}
}
fn module_type_from_name(&self, module_name: &ModuleName) -> Option<Type<'db>> {
resolve_module(self.db(), module_name)
.map(|module| Type::module_literal(self.db(), self.file(), &module))
@ -5775,13 +5889,15 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> {
let current_file = self.file();
let mut is_nonlocal_binding = false;
if let Some(name) = expr.as_name() {
let skip_non_global_scopes = place_table
.place_id_by_name(name)
.is_some_and(|symbol_id| self.skip_non_global_scopes(file_scope_id, symbol_id));
if skip_non_global_scopes {
return global_symbol(self.db(), self.file(), name);
if let Some(symbol_id) = place_table.place_id_by_name(name) {
if self.skip_non_global_scopes(file_scope_id, symbol_id) {
return global_symbol(self.db(), self.file(), name);
}
is_nonlocal_binding = self
.index
.symbol_is_nonlocal_in_scope(symbol_id, file_scope_id);
}
}
@ -5794,7 +5910,7 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> {
// a local variable or not in function-like scopes. If a variable has any bindings in a
// function-like scope, it is considered a local variable; it never references another
// scope. (At runtime, it would use the `LOAD_FAST` opcode.)
if has_bindings_in_this_scope && scope.is_function_like(db) {
if has_bindings_in_this_scope && scope.is_function_like(db) && !is_nonlocal_binding {
return Place::Unbound.into();
}