[ty] more cases for the class body global fallback

crates/ty_python_semantic/src/semantic_index/symbol.rs

@@ -80,6 +80,38 @@ impl Symbol {
         self.flags.contains(SymbolFlags::MARKED_NONLOCAL)
     }
 
+    /// Is the symbol defined in this scope, vs referring to some enclosing scope?
+    ///
+    /// There are three common cases where a name refers to an enclosing scope:
+    ///
+    /// 1. explicit `global` variables
+    /// 2. explicit `nonlocal` variables
+    /// 3. "free" variables, which are used in a scope where they're neither bound nor declared
+    ///
+    /// Note that even if `is_local` is false, that doesn't necessarily mean there's an enclosing
+    /// scope that resolves the reference. The symbol could be a built-in like `print`, or a name
+    /// error at runtime, or a global variable added dynamically with e.g. `globals()`.
+    ///
+    /// XXX: There's a fourth case that we don't (can't) handle here. A variable that's bound or
+    /// declared (anywhere) in a class body, but used before it's bound (at runtime), resolves
+    /// (unbelievably) to the global scope. For example:
+    /// ```py
+    /// x = 42
+    /// def f():
+    ///     x = 43
+    ///     class Foo:
+    ///         print(x)  # 42 (never 43)
+    ///         if secrets.randbelow(2):
+    ///             x = 44
+    ///         print(x)  # 42 or 44
+    /// ```
+    /// In cases like this, the resolution isn't known until runtime, and in fact it varies from
+    /// one use to the next. The semantic index alone can't resolve this, and instead it's a
+    /// special case in type inference (see `infer_place_load`).
+    pub(crate) fn is_local(&self) -> bool {
+        !self.is_global() && !self.is_nonlocal() && (self.is_bound() || self.is_declared())
+    }
+
     pub(crate) const fn is_reassigned(&self) -> bool {
         self.flags.contains(SymbolFlags::IS_REASSIGNED)
     }

crates/ty_python_semantic/src/types/infer.rs

@@ -6664,30 +6664,19 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> {
         if let Some(use_id) = use_id {
             constraint_keys.push((file_scope_id, ConstraintKey::UseId(use_id)));
         }
-        let local_is_unbound = local_scope_place.is_unbound();
 
         let place = PlaceAndQualifiers::from(local_scope_place).or_fall_back_to(db, || {
-            let has_bindings_in_this_scope = match place_table.place_id(place_expr) {
-                Some(place_id) => place_table.place(place_id).is_bound(),
-                None => {
-                    assert!(
-                        self.deferred_state.in_string_annotation(),
-                        "Expected the place table to create a place for every Name node"
-                    );
-                    false
-                }
-            };
-
             let current_file = self.file();
-
-            let mut is_nonlocal_binding = false;
+            let mut symbol_resolves_locally = false;
             if let Some(symbol) = place_expr.as_symbol() {
                 if let Some(symbol_id) = place_table.symbol_id(symbol.name()) {
-                    // If we try to access a variable in a class before it has been defined,
-                    // the lookup will fall back to global.
-                    let fallback_to_global = local_is_unbound
-                        && has_bindings_in_this_scope
-                        && scope.node(db).scope_kind().is_class();
+                    // Footgun: `place_expr` and `symbol` were probably constructed with all-zero
+                    // flags. We need to read the place table to get correct flags.
+                    symbol_resolves_locally = place_table.symbol(symbol_id).is_local();
+                    // If we try to access a variable in a class before it has been defined, the
+                    // lookup will fall back to global. See the comment on `Symbol::is_local`.
+                    let fallback_to_global =
+                        scope.node(db).scope_kind().is_class() && symbol_resolves_locally;
                     if self.skip_non_global_scopes(file_scope_id, symbol_id) || fallback_to_global {
                         return global_symbol(self.db(), self.file(), symbol.name()).map_type(
                             |ty| {
@@ -6699,22 +6688,15 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> {
                             },
                         );
                     }
-                    is_nonlocal_binding = self
-                        .index
-                        .symbol_is_nonlocal_in_scope(symbol_id, file_scope_id);
                 }
             }
 
-            // If it's a function-like scope and there is one or more binding in this scope (but
-            // none of those bindings are visible from where we are in the control flow), we cannot
-            // fallback to any bindings in enclosing scopes. As such, we can immediately short-circuit
-            // here and return `Place::Unbound`.
-            //
-            // This is because Python is very strict in its categorisation of whether a variable is
-            // 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) && !is_nonlocal_binding {
+            // Symbols that are bound or declared in the local scope, and not marked `nonlocal` or
+            // `global`, never refer to an enclosing scope. (If you reference such a symbol before
+            // it's bound, you get an `UnboundLocalError`.) Short-circuit instead of walking
+            // enclosing scopes in this case. The one exception to this rule is the global fallback
+            // in class bodies, which we already handled above.
+            if symbol_resolves_locally {
                 return Place::Unbound.into();
             }