don't do nested bindings inference if we have a declared type

crates/ty_python_semantic/resources/mdtest/public_types.md

@@ -312,7 +312,7 @@ def outer() -> None:
     set_x()
 
     def inner() -> None:
-        reveal_type(x)  # revealed: None | Literal[1]
+        reveal_type(x)  # revealed: Literal[1] | None
     inner()
 ```
 

crates/ty_python_semantic/resources/mdtest/scopes/global.md

@@ -201,7 +201,7 @@ x = 42
 
 def f():
     global x
-    reveal_type(x)  # revealed: Unknown | Literal[42, "56"]
+    reveal_type(x)  # revealed: Unknown | Literal["56", 42]
     x = "56"
     reveal_type(x)  # revealed: Literal["56"]
 ```

crates/ty_python_semantic/resources/mdtest/scopes/nonlocal.md

@@ -106,12 +106,12 @@ def a():
                 nonlocal x
                 # It's counterintuitive that 4 gets included here, since we haven't reached the
                 # binding in this scope, but this function might get called more than once.
-                reveal_type(x)  # revealed: Literal[2, 3, 4]
+                reveal_type(x)  # revealed: Literal[3, 4, 2]
                 x = 4
                 reveal_type(x)  # revealed: Literal[4]
 
                 def e():
-                    reveal_type(x)  # revealed: Literal[2, 3, 4]
+                    reveal_type(x)  # revealed: Literal[3, 4, 2]
 ```
 
 In addition to parent scopes, we also consider sibling scopes, child scopes,
@@ -131,7 +131,7 @@ def a():
         def d():
             nonlocal x
             x = 3
-        reveal_type(x)  # revealed: Literal[1, 2, 3]
+        reveal_type(x)  # revealed: Literal[2, 3, 1]
     # `x` is local here, so we don't look at nested scopes.
     reveal_type(x)  # revealed: Literal[1]
 ```

crates/ty_python_semantic/src/place.rs

@@ -651,32 +651,6 @@ fn place_by_id<'db>(
 ) -> PlaceAndQualifiers<'db> {
     let use_def = use_def_map(db, scope);
 
-    // If there are any nested bindings (via `global` or `nonlocal` variables) for this symbol,
-    // infer them and union the results. Nested bindings aren't allowed to have declarations or
-    // qualifiers, and we can just union their inferred types.
-    let mut nested_bindings_union = UnionBuilder::new(db);
-    if let Some(symbol_id) = place_id.as_symbol() {
-        let current_place_table = place_table(db, scope);
-        let symbol = current_place_table.symbol(symbol_id);
-        for nested_file_scope_id in place_table(db, scope).nested_scopes_with_bindings(symbol_id) {
-            let nested_scope_id = nested_file_scope_id.to_scope_id(db, scope.file(db));
-            let nested_place_table = place_table(db, nested_scope_id);
-            let nested_symbol_id = nested_place_table
-                .symbol_id(symbol.name())
-                .expect("nested_scopes_with_bindings says this reference exists");
-            let nested_place = place_by_id(
-                db,
-                nested_scope_id,
-                ScopedPlaceId::Symbol(nested_symbol_id),
-                RequiresExplicitReExport::No,
-                ConsideredDefinitions::AllReachable,
-            );
-            if let Place::Type(nested_type, _) = nested_place.place {
-                nested_bindings_union.add_in_place(nested_type);
-            }
-        }
-    }
-
     // If the place is declared, the public type is based on declarations; otherwise, it's based
     // on inference from bindings.
 
@@ -692,6 +666,47 @@ fn place_by_id<'db>(
         ConsideredDefinitions::AllReachable => use_def.all_reachable_bindings(place_id),
     };
 
+    // If there are any nested bindings (via `global` or `nonlocal` variables) for this symbol,
+    // infer them and union the results. Note that this is potentially recursive, and we can
+    // trigger fixed-point iteration here in cases like this:
+    // ```
+    // def f():
+    //     x = 1
+    //     def g():
+    //         nonlocal x
+    //         x += 1
+    // ```
+    let nested_bindings = || {
+        // For performance reasons, avoid creating a union unless we have more one binding.
+        let mut union = UnionBuilder::new(db);
+        if let Some(symbol_id) = place_id.as_symbol() {
+            let current_place_table = place_table(db, scope);
+            let symbol = current_place_table.symbol(symbol_id);
+            for &nested_file_scope_id in
+                place_table(db, scope).nested_scopes_with_bindings(symbol_id)
+            {
+                let nested_scope_id = nested_file_scope_id.to_scope_id(db, scope.file(db));
+                let nested_place_table = place_table(db, nested_scope_id);
+                let nested_symbol_id = nested_place_table
+                    .symbol_id(symbol.name())
+                    .expect("nested_scopes_with_bindings says this reference exists");
+                let place = place_by_id(
+                    db,
+                    nested_scope_id,
+                    ScopedPlaceId::Symbol(nested_symbol_id),
+                    RequiresExplicitReExport::No,
+                    ConsideredDefinitions::AllReachable,
+                );
+                // Nested bindings aren't allowed to have declarations or qualifiers, so we can
+                // just extract their inferred types.
+                if let Place::Type(nested_type, _) = place.place {
+                    union.add_in_place(nested_type);
+                }
+            }
+        }
+        union
+    };
+
     // If a symbol is undeclared, but qualified with `typing.Final`, we use the right-hand side
     // inferred type, without unioning with `Unknown`, because it can not be modified.
     if let Some(qualifiers) = declared
@@ -751,17 +766,11 @@ fn place_by_id<'db>(
                     // TODO: We probably don't want to report `Bound` here. This requires a bit of
                     // design work though as we might want a different behavior for stubs and for
                     // normal modules.
-                    Place::Type(
-                        nested_bindings_union.add(declared_ty).build(),
-                        Boundness::Bound,
-                    )
+                    Place::Type(nested_bindings().add(declared_ty).build(), Boundness::Bound)
                 }
                 // Place is possibly undeclared and (possibly) bound
                 Place::Type(inferred_ty, boundness) => Place::Type(
-                    nested_bindings_union
-                        .add(inferred_ty)
-                        .add(declared_ty)
-                        .build(),
+                    nested_bindings().add(inferred_ty).add(declared_ty).build(),
                     if boundness_analysis == BoundnessAnalysis::AssumeBound {
                         Boundness::Bound
                     } else {
@@ -784,13 +793,12 @@ fn place_by_id<'db>(
 
             // If there are nested bindings, union whatever we inferred from those into what we've
             // inferred here.
-            if let Some(nested_bindings_type) = nested_bindings_union.try_build() {
-                match &mut inferred {
-                    Place::Type(inferred_type, _) => {
-                        *inferred_type =
-                            UnionType::from_elements(db, [*inferred_type, nested_bindings_type]);
-                    }
-                    Place::Unbound => {
+            match &mut inferred {
+                Place::Type(inferred_type, _) => {
+                    *inferred_type = nested_bindings().add(*inferred_type).build();
+                }
+                Place::Unbound => {
+                    if let Some(nested_bindings_type) = nested_bindings().try_build() {
                         inferred = Place::Type(nested_bindings_type, Boundness::PossiblyUnbound);
                     }
                 }

crates/ty_python_semantic/src/semantic_index/builder.rs

@@ -14,7 +14,7 @@ use ruff_python_ast::{self as ast, NodeIndex, PySourceType, PythonVersion};
 use ruff_python_parser::semantic_errors::{
     SemanticSyntaxChecker, SemanticSyntaxContext, SemanticSyntaxError, SemanticSyntaxErrorKind,
 };
-use ruff_text_size::{Ranged, TextRange};
+use ruff_text_size::TextRange;
 
 use crate::ast_node_ref::AstNodeRef;
 use crate::module_name::ModuleName;
@@ -66,40 +66,42 @@ impl Loop {
     }
 }
 
-struct ScopeInfo<'ast> {
+struct ScopeInfo {
     file_scope_id: FileScopeId,
 
     /// Current loop state; None if we are not currently visiting a loop
     current_loop: Option<Loop>,
 
-    /// Symbols from scopes nested inside of this one that haven't yet been resolved to a
-    /// definition. They might end up resolving in this scope, or in an enclosing scope.
+    /// `nonlocal` variables from scopes nested inside of this one that haven't yet been resolved
+    /// to a definition. They might end up resolving in this scope, or in an enclosing scope.
     ///
-    /// When we pop scopes, we merge any unresolved free variables into the parent scope's
-    /// collection. The reason we need to collect free variables for each scope separately, instead
-    /// of just having one map for the whole builder, is because of sibling scope arrangements like
-    /// this:
+    /// When we pop scopes, we merge any unresolved nonlocals into the parent scope's collection.
+    /// The reason we need to track them for each scope separately, instead of using one map for
+    /// the whole builder, is because of sibling scope arrangements like this:
     /// ```py
     /// def f():
     ///     def g():
-    ///         # When we pop `g`, this `x` goes in `f`'s set of free variables.
+    ///         # When we pop `g`, this `x` goes in `f`'s set of unresolved nonlocals.
     ///         nonlocal x
     ///     def h():
-    ///         # When we pop `h`, this binding of `x` won't resolve the free variable from `g`,
-    ///         # because it's not in `h`'s set of free variables.
+    ///         # When we pop `h`, this binding of `x` will *not* resolve the nonlocal from `g`,
+    ///         # because it's not in `h`'s set of unresolved nonlocals.
     ///         x = 1
-    ///     # When we pop `f`, this binding of `x` will resolve the free variable from `g`.
+    ///     # When we pop `f`, this binding of `x` will resolve the nonlocal from `g`.
     ///     x = 1
     /// ```
-    free_variables: FxHashMap<ast::name::Name, Vec<FreeVariable<'ast>>>,
+    ///
+    /// Currently we only track explicit nonlocals, because ordinary "free" variables referring to
+    /// enclosing scopes can't be bound and can't trigger semantic syntax errors. Callers who need
+    /// to resolve free variables (e.g. `infer_place_load`) need to walk parent scopes until they
+    /// find one where `Symbol::is_local` or `Symbol::is_global` is true. If we wanted to
+    /// pre-record more of that information, we could expand this.
+    unresolved_nonlocals: FxHashMap<ast::name::Name, Vec<UnresolvedNonlocal>>,
 }
 
-struct FreeVariable<'ast> {
+struct UnresolvedNonlocal {
     scope_id: FileScopeId,
-    // If this variable is `nonlocal`, then this is `Some` reference to its identifier in the
-    // `nonlocal` statement. In that case, it's an error if we don't resolve it before we reach the
-    // global scope (or if we resolve it in a scope where it's `global`).
-    nonlocal_identifier: Option<&'ast ast::Identifier>,
+    range: TextRange,
 }
 
 pub(super) struct SemanticIndexBuilder<'db, 'ast> {
@@ -108,7 +110,7 @@ pub(super) struct SemanticIndexBuilder<'db, 'ast> {
     file: File,
     source_type: PySourceType,
     module: &'ast ParsedModuleRef,
-    scope_stack: Vec<ScopeInfo<'ast>>,
+    scope_stack: Vec<ScopeInfo>,
     /// The assignments we're currently visiting, with
     /// the most recent visit at the end of the Vec
     current_assignments: Vec<CurrentAssignment<'ast, 'db>>,
@@ -197,13 +199,13 @@ impl<'db, 'ast> SemanticIndexBuilder<'db, 'ast> {
         builder
     }
 
-    fn current_scope_info(&self) -> &ScopeInfo<'ast> {
+    fn current_scope_info(&self) -> &ScopeInfo {
         self.scope_stack
             .last()
             .expect("SemanticIndexBuilder should have created a root scope")
     }
 
-    fn current_scope_info_mut(&mut self) -> &mut ScopeInfo<'ast> {
+    fn current_scope_info_mut(&mut self) -> &mut ScopeInfo {
         self.scope_stack
             .last_mut()
             .expect("SemanticIndexBuilder should have created a root scope")
@@ -305,7 +307,7 @@ impl<'db, 'ast> SemanticIndexBuilder<'db, 'ast> {
         self.scope_stack.push(ScopeInfo {
             file_scope_id,
             current_loop: None,
-            free_variables: FxHashMap::default(),
+            unresolved_nonlocals: FxHashMap::default(),
         });
     }
 
@@ -477,7 +479,7 @@ impl<'db, 'ast> SemanticIndexBuilder<'db, 'ast> {
 
         let ScopeInfo {
             file_scope_id: popped_scope_id,
-            free_variables: mut popped_free_variables,
+            unresolved_nonlocals: mut popped_unresolved_nonlocals,
             ..
         } = self
             .scope_stack
@@ -499,145 +501,93 @@ impl<'db, 'ast> SemanticIndexBuilder<'db, 'ast> {
             self.record_lazy_snapshots(popped_scope_id);
         }
 
-        // If we've popped a scope that free variables from nested (previously popped) scopes can
-        // refer to (i.e. not a class body), try to resolve outstanding free variables.
+        // If we've popped a scope that nonlocals from nested (previously popped) scopes can refer
+        // to (i.e. not a class body), try to resolve them.
         if kind.is_function_like() || popped_scope_id.is_global() {
-            // Look up each free variable name in the popped scope, and see if we've resolved it.
-            // Collect these in a separate list, to avoid borrowck woes.
-            struct Resolution {
-                name: ast::name::Name,
-                symbol_id: ScopedSymbolId,
-                // Either the symbol is declared `global`, or this is the global scope.
-                is_global: bool,
-            }
-            let mut resolutions = Vec::new();
-            for name in popped_free_variables.keys() {
-                if let Some(symbol_id) = self.place_tables[popped_scope_id].symbol_id(name.as_str())
-                {
-                    // If a name is local or `global` here (i.e. bound or declared, and not marked
-                    // `nonlocal`), then free variables of that name resolve here. Note that
-                    // popping scopes in the normal stack order means that free variables resolve
-                    // (correctly) to the closest scope with a matching definition.
-                    let symbol = self.place_tables[popped_scope_id].symbol(symbol_id);
-                    if symbol.is_local() || symbol.is_global() {
-                        resolutions.push(Resolution {
-                            name: name.clone(),
-                            symbol_id,
-                            is_global: symbol.is_global() || popped_scope_id.is_global(),
-                        });
-                    }
-                }
-            }
-
-            // Remove each resolved name along with all its references from
-            // `popped_free_variables`. For each reference, if it's bound in its nested scope, add
-            // an entry to `nested_scopes_with_bindings` in the popped scope's symbol table. This
-            // is also where we flag any `nonlocal` statements that resolve to globals, which is a
-            // semantic syntax error.
-            for resolution in resolutions {
-                let resolved_variables = popped_free_variables.remove(&resolution.name).unwrap();
-                for FreeVariable {
-                    scope_id: nested_scope_id,
-                    nonlocal_identifier,
-                } in resolved_variables
-                {
-                    let nested_symbol_is_nonlocal = nonlocal_identifier.is_some();
-                    if nested_symbol_is_nonlocal && resolution.is_global {
-                        // If the symbol is declared `nonlocal` in the nested scope (rather than
-                        // just used without a local binding or declaration), then it's a syntax
-                        // error for it to resolve to the global scope or to a `global` statement.
-                        self.report_semantic_error(SemanticSyntaxError {
-                            kind: SemanticSyntaxErrorKind::NoBindingForNonlocal(
-                                resolution.name.clone().into(),
-                            ),
-                            range: nonlocal_identifier.unwrap().range(),
-                            python_version: self.python_version,
-                        });
-                    } else {
-                        let nested_place_table = &self.place_tables[nested_scope_id];
-                        let nested_symbol_id =
-                            nested_place_table.symbol_id(&resolution.name).unwrap();
-                        let nested_symbol = nested_place_table.symbol(nested_symbol_id);
-                        if nested_symbol.is_bound() {
-                            self.place_tables[popped_scope_id].add_nested_scope_with_binding(
-                                resolution.symbol_id,
-                                nested_scope_id,
-                            );
+            popped_unresolved_nonlocals.retain(|name, nonlocals_with_this_name| {
+                let popped_place_table = &self.place_tables[popped_scope_id];
+                if let Some(symbol_id) = popped_place_table.symbol_id(name.as_str()) {
+                    let symbol = popped_place_table.symbol(symbol_id);
+                    let symbol_is_resolved = symbol.is_local() || symbol.is_global();
+                    let resolution_is_global = symbol.is_global() || popped_scope_id.is_global();
+                    if symbol_is_resolved {
+                        for &mut UnresolvedNonlocal {
+                            scope_id: nested_scope_id,
+                            range,
+                        } in nonlocals_with_this_name
+                        {
+                            if resolution_is_global {
+                                // It's a syntax error for a nonlocal variable to resolve to the
+                                // global scope or to a `global` statement in an enclosing scope.
+                                self.report_semantic_error(SemanticSyntaxError {
+                                    kind: SemanticSyntaxErrorKind::NoBindingForNonlocal(
+                                        name.clone().into(),
+                                    ),
+                                    range,
+                                    python_version: self.python_version,
+                                });
+                                continue;
+                            }
+                            let nested_place_table = &self.place_tables[nested_scope_id];
+                            let nested_symbol_id = nested_place_table.symbol_id(name).unwrap();
+                            let nested_symbol = nested_place_table.symbol(nested_symbol_id);
+                            if nested_symbol.is_bound() {
+                                self.place_tables[popped_scope_id]
+                                    .add_nested_scope_with_binding(symbol_id, nested_scope_id);
+                            }
                         }
+                        // This name was resolved. Remove it and its references.
+                        return false;
                     }
                 }
-            }
+                // This name was not resolved. Retain it. We'll add it to the parent scope's
+                // collection below.
+                true
+            });
         }
 
         if popped_scope_id.is_global() {
-            // If we've popped the global/module scope, any remaining free variables are
-            // unresolved. The common case for these is built-ins like `print`, and rarer cases are
-            // things like direct insertions into `globals()`. However, if any `nonlocal` free
-            // variables are still unresolved, that's another syntax error.
+            // If we've popped the global/module scope, still-unresolved `nonlocal` variables are
+            // another syntax error.
             debug_assert!(self.scope_stack.is_empty());
-            for (name, variables) in &popped_free_variables {
-                for variable in variables {
-                    if let Some(nonlocal_identifier) = variable.nonlocal_identifier {
-                        self.report_semantic_error(SemanticSyntaxError {
-                            kind: SemanticSyntaxErrorKind::NoBindingForNonlocal(
-                                name.clone().into(),
-                            ),
-                            range: nonlocal_identifier.range(),
-                            python_version: self.python_version,
-                        });
-                    }
+            for (name, nonlocals) in &popped_unresolved_nonlocals {
+                for nonlocal in nonlocals {
+                    self.report_semantic_error(SemanticSyntaxError {
+                        kind: SemanticSyntaxErrorKind::NoBindingForNonlocal(name.clone().into()),
+                        range: nonlocal.range,
+                        python_version: self.python_version,
+                    });
                 }
             }
         } else {
-            // Otherwise, add any still-unresolved free variables from nested scopes to the parent
-            // scope's collection, and walk the popped scope's symbol table to collect any new free
-            // variables. During that walk, also record references to global variables.
-            let parent_free_variables = &mut self
+            // Otherwise, add any still-unresolved nonlocals from nested scopes to the parent
+            // scope's collection.
+            let parent_unresolved_nonlocals = &mut self
                 .scope_stack
                 .last_mut() // current_scope_info_mut() would be a borrock error here
                 .expect("this is not the global/module scope")
-                .free_variables;
-            for (name, variables) in popped_free_variables {
-                parent_free_variables
+                .unresolved_nonlocals;
+            for (name, variables) in popped_unresolved_nonlocals {
+                parent_unresolved_nonlocals
                     .entry(name)
                     .or_default()
                     .extend(variables);
             }
-            let popped_place_table = &self.place_tables[popped_scope_id];
-            let mut bound_global_symbols = Vec::new();
-            for symbol in popped_place_table.symbols() {
-                // Collect new implicit (not `nonlocal`) free variables.
-                //
-                // NOTE: Because these variables aren't bound (won't wind up in
-                // `nested_scopes_with_bindings`) and aren't `nonlocal` (can't trigger `nonlocal`
-                // syntax errors), collecting them currently has no effect. We could consider
-                // removing this bit and renaming `free_variables` to say `unresolved_nonlocals`?
-                if symbol.is_used()
-                    && !symbol.is_bound()
-                    && !symbol.is_declared()
-                    && !symbol.is_global()
-                    // `nonlocal` variables are handled in `visit_stmt`, which lets us stash an AST
-                    // reference.
-                    && !symbol.is_nonlocal()
-                {
-                    parent_free_variables
-                        .entry(symbol.name().clone())
-                        .or_default()
-                        .push(FreeVariable {
-                            scope_id: popped_scope_id,
-                            nonlocal_identifier: None,
-                        });
-                }
 
-                // Record bindings of global variables. Put these in a temporary Vec as another
-                // borrowck workaround.
+            // Also, update the global/module symbol table with any bound `global` variables in
+            // this scope. We do this here, rather than when we visit the `global` statement,
+            // because at that point we don't know whether the variable is bound.
+            let mut bound_global_symbols = Vec::new();
+            for symbol in self.place_tables[popped_scope_id].symbols() {
+                // Record bindings of global variables in a temporary Vec as a borrowck workaround.
+                // We could get clever here, but `global` variables are relatively rare, and
+                // allocating a small Vec in those cases isn't expensive.
                 if symbol.is_global() && symbol.is_bound() {
                     bound_global_symbols.push(symbol.name().clone());
                 }
             }
-
-            // Update the global scope with those references to globals, now that
-            // `popped_place_table` and `parent_free_variables` are no longer borrowed.
+            // Update the global scope with those bindings, now that `self.place_tables` is no
+            // longer borrowed.
             for symbol_name in bound_global_symbols {
                 // Add this symbol to the global scope, if it isn't there already.
                 let global_symbol_id = self.add_symbol_to_scope(symbol_name, FileScopeId::global());
@@ -2313,9 +2263,9 @@ impl<'ast> Visitor<'ast> for SemanticIndexBuilder<'_, 'ast> {
                     self.current_place_table_mut()
                         .symbol_mut(symbol_id)
                         .mark_global();
-                    // We'll add this symbol to the global scope in `pop_scope`, at the same time
-                    // we're collecting free variables. That lets us record whether it's bound in
-                    // this scope, which we don't know yet.
+                    // We'll add this symbol to the global scope in `pop_scope`, after we resolve
+                    // nonlocals. That lets us record whether it's bound in this scope, which we
+                    // don't know yet.
                 }
                 walk_stmt(self, stmt);
             }
@@ -2362,20 +2312,20 @@ impl<'ast> Visitor<'ast> for SemanticIndexBuilder<'_, 'ast> {
                     self.current_place_table_mut()
                         .symbol_mut(symbol_id)
                         .mark_nonlocal();
-                    // Add this symbol to the parent scope's set of free variables. (It would also
-                    // work to add it to this scope's set, which will get folded into the parent's
-                    // in `pop_scope`. But since it can't possibly resolve here, we might as well
-                    // spare an allocation.) We checked above that we aren't in the module scope,
-                    // so there's definitely a parent scope.
+                    // Add this symbol to the parent scope's set of unresolved nonlocals. (It would
+                    // also work to add it to this scope's set, which will get folded into the
+                    // parent's in `pop_scope`. But since it can't possibly resolve here, we might
+                    // as well try to spare an allocation.) We checked above that we aren't in the
+                    // module scope, so there's definitely a parent scope.
                     let parent_scope_index = self.scope_stack.len() - 2;
                     let parent_scope_info = &mut self.scope_stack[parent_scope_index];
                     parent_scope_info
-                        .free_variables
+                        .unresolved_nonlocals
                         .entry(name.id.clone())
                         .or_default()
-                        .push(FreeVariable {
+                        .push(UnresolvedNonlocal {
                             scope_id,
-                            nonlocal_identifier: Some(name),
+                            range: name.range,
                         });
                 }
                 walk_stmt(self, stmt);

crates/ty_python_semantic/src/semantic_index/symbol.rs

@@ -276,6 +276,9 @@ impl SymbolTableBuilder {
             .map
             .shrink_to_fit(|id| SymbolTable::hash_name(&table.symbols[*id].name));
         table.nested_scopes_with_bindings.shrink_to_fit();
+        for scopes_vec in table.nested_scopes_with_bindings.values_mut() {
+            scopes_vec.shrink_to_fit();
+        }
         table
     }
 }