fix snapshot behavior in annotation scopes

annotation scopes can see names defined in an immediately-enclosing class scope

Author: mtshiba

crates/ty_python_semantic/resources/mdtest/annotations/deferred.md

@@ -87,9 +87,23 @@ class Foo:
         class Baz[T: Foo]:
             pass
 
+        # error: [unresolved-reference] "Name `Foo` used when not defined"
+        # error: [unresolved-reference] "Name `Bar` used when not defined"
+        class Qux(Foo, Bar, Baz):
+            pass
+
+        # error: [unresolved-reference] "Name `Foo` used when not defined"
+        # error: [unresolved-reference] "Name `Bar` used when not defined"
+        class Quux[_T](Foo, Bar, Baz):
+            pass
+
         # error: [unresolved-reference]
         type S = a
         type T = b
+        type U = Foo
+        # error: [unresolved-reference]
+        type V = Bar
+        type W = Baz
 
     def h[T: Bar]():
         # error: [unresolved-reference]
@@ -141,9 +155,23 @@ class Foo:
         class Baz[T: Foo]:
             pass
 
+        # error: [unresolved-reference] "Name `Foo` used when not defined"
+        # error: [unresolved-reference] "Name `Bar` used when not defined"
+        class Qux(Foo, Bar, Baz):
+            pass
+
+        # error: [unresolved-reference] "Name `Foo` used when not defined"
+        # error: [unresolved-reference] "Name `Bar` used when not defined"
+        class Quux[_T](Foo, Bar, Baz):
+            pass
+
         # error: [unresolved-reference]
         type S = a
         type T = b
+        type U = Foo
+        # error: [unresolved-reference]
+        type V = Bar
+        type W = Baz
 
     def h[T: Bar]():
         # error: [unresolved-reference]

crates/ty_python_semantic/src/semantic_index/builder.rs

@@ -314,6 +314,8 @@ impl<'db, 'ast> SemanticIndexBuilder<'db, 'ast> {
 
     // Records snapshots of the place states visible from the current eager scope.
     fn record_eager_snapshots(&mut self, popped_scope_id: FileScopeId) {
+        let popped_scope = &self.scopes[popped_scope_id];
+
         // If the scope that we just popped off is an eager scope, we need to "lock" our view of
         // which bindings reach each of the uses in the scope. Loop through each enclosing scope,
         // looking for any that bind each place.
@@ -350,11 +352,14 @@ impl<'db, 'ast> SemanticIndexBuilder<'db, 'ast> {
                     nested_scope: popped_scope_id,
                     nested_laziness: ScopeLaziness::Eager,
                 };
+                let is_immediately_enclosing_scope = popped_scope.kind().is_annotation()
+                    && popped_scope.parent() == Some(enclosing_scope_id);
                 let eager_snapshot = self.use_def_maps[enclosing_scope_id]
                     .snapshot_enclosing_state(
                         enclosing_place_id,
                         enclosing_scope_kind,
                         enclosing_place,
+                        is_immediately_enclosing_scope,
                     );
                 self.enclosing_snapshots.insert(key, eager_snapshot);
             }
@@ -429,6 +434,7 @@ impl<'db, 'ast> SemanticIndexBuilder<'db, 'ast> {
                     enclosed_symbol_id.into(),
                     enclosing_scope_kind,
                     enclosing_place.into(),
+                    false,
                 );
                 self.enclosing_snapshots.insert(key, lazy_snapshot);
             }

crates/ty_python_semantic/src/semantic_index/use_def.rs

@@ -1186,17 +1186,22 @@ impl<'db> UseDefMapBuilder<'db> {
     pub(super) fn snapshot_enclosing_state(
         &mut self,
         enclosing_place: ScopedPlaceId,
-        scope: ScopeKind,
+        enclosing_scope: ScopeKind,
         enclosing_place_expr: PlaceExprRef,
+        is_immediately_enclosing_scope: bool,
     ) -> ScopedEnclosingSnapshotId {
         let bindings = match enclosing_place {
             ScopedPlaceId::Symbol(symbol) => self.symbol_states[symbol].bindings(),
             ScopedPlaceId::Member(member) => self.member_states[member].bindings(),
         };
 
+        let is_class_symbol = enclosing_scope.is_class() && enclosing_place.is_symbol();
         // Names bound in class scopes are never visible to nested scopes (but attributes/subscripts are visible),
         // so we never need to save eager scope bindings in a class scope.
-        if (scope.is_class() && enclosing_place.is_symbol()) || !enclosing_place_expr.is_bound() {
+        // There is one exception to this rule: annotation scopes can see
+        // names defined in an immediately-enclosing class scope.
+        if (is_class_symbol && !is_immediately_enclosing_scope) || !enclosing_place_expr.is_bound()
+        {
             self.enclosing_snapshots.push(EnclosingSnapshot::Constraint(
                 bindings.unbound_narrowing_constraint(),
             ))