[ty] improve lazy scope place lookup (#19321)

Co-authored-by: David Peter <sharkdp@users.noreply.github.com>
Co-authored-by: Carl Meyer <carl@oddbird.net>
Co-authored-by: Micha Reiser <micha@reiser.io>

crates/ty_python_semantic/resources/mdtest/del.md

@@ -123,7 +123,7 @@ def enclosing():
         nonlocal x
         def bar():
             # allowed, refers to `x` in `enclosing`
-            reveal_type(x)  # revealed: Unknown | Literal[2]
+            reveal_type(x)  # revealed: Literal[2]
         bar()
         del x  # allowed, deletes `x` in `enclosing` (though we don't track that)
 ```

crates/ty_python_semantic/resources/mdtest/narrow/conditionals/nested.md

@@ -238,6 +238,69 @@ def f(x: str | None):
 
         [reveal_type(x) for _ in range(1)]  # revealed: str
 
+    # When there is a reassignment, any narrowing constraints on the place are invalidated in lazy scopes.
+    x = None
+```
+
+If a variable defined in a private scope is never reassigned, narrowing remains in effect in the
+inner lazy scope.
+
+```py
+def f(const: str | None):
+    if const is not None:
+        def _():
+            # The `const is not None` narrowing constraint is still valid since `const` has not been reassigned
+            reveal_type(const)  # revealed: str
+
+        class C2:
+            reveal_type(const)  # revealed: str
+
+        [reveal_type(const) for _ in range(1)]  # revealed: str
+```
+
+And even if there is an attribute or subscript assignment to the variable, narrowing of the variable
+is still valid in the inner lazy scope.
+
+```py
+def f(l: list[str | None] | None):
+    if l is not None:
+        def _():
+            reveal_type(l)  # revealed: list[str | None]
+        l[0] = None
+
+def f(a: A):
+    if a:
+        def _():
+            reveal_type(a)  # revealed: A & ~AlwaysFalsy
+    a.x = None
+```
+
+Narrowing is invalidated if a `nonlocal` declaration is made within a lazy scope.
+
+```py
+def f(non_local: str | None):
+    if non_local is not None:
+        def _():
+            nonlocal non_local
+            non_local = None
+
+        def _():
+            reveal_type(non_local)  # revealed: str | None
+
+def f(non_local: str | None):
+    def _():
+        nonlocal non_local
+        non_local = None
+    if non_local is not None:
+        def _():
+            reveal_type(non_local)  # revealed: str | None
+```
+
+The same goes for public variables, attributes, and subscripts, because it is difficult to track all
+of their changes.
+
+```py
+def f():
     if g is not None:
         def _():
             reveal_type(g)  # revealed: str | None
@@ -249,6 +312,7 @@ def f(x: str | None):
 
     if a.x is not None:
         def _():
+            # Lazy nested scope narrowing is not performed on attributes/subscripts because it's difficult to track their changes.
             reveal_type(a.x)  # revealed: Unknown | str | None
 
         class D:
@@ -282,7 +346,7 @@ l: list[str | Literal[1] | None] = [None]
 
 def f(x: str | Literal[1] | None):
     class C:
-        if x is not None:
+        if x is not None:  # TODO: should be an unresolved-reference error
             def _():
                 if x != 1:
                     reveal_type(x)  # revealed: str | None
@@ -293,6 +357,38 @@ def f(x: str | Literal[1] | None):
 
             [reveal_type(x) for _ in range(1) if x != 1]  # revealed: str
 
+        x = None
+
+    def _():
+        # error: [unresolved-reference]
+        if x is not None:
+            def _():
+                if x != 1:
+                    reveal_type(x)  # revealed: Never
+        x = None
+
+def f(const: str | Literal[1] | None):
+    class C:
+        if const is not None:
+            def _():
+                if const != 1:
+                    # TODO: should be `str`
+                    reveal_type(const)  # revealed: str | None
+
+            class D:
+                if const != 1:
+                    reveal_type(const)  # revealed: str
+
+            [reveal_type(const) for _ in range(1) if const != 1]  # revealed: str
+
+    def _():
+        if const is not None:
+            def _():
+                if const != 1:
+                    reveal_type(const)  # revealed: str
+
+def f():
+    class C:
         if g is not None:
             def _():
                 if g != 1:

crates/ty_python_semantic/resources/mdtest/public_types.md

@@ -20,10 +20,7 @@ def outer() -> None:
     x = A()
 
     def inner() -> None:
-        # TODO: We might ideally be able to eliminate `Unknown` from the union here since `x` resolves to an
+        reveal_type(x)  # revealed: A | B
-        # outer scope that is a function scope (as opposed to module global scope), and `x` is never declared
-        # nonlocal in a nested scope that also assigns to it.
-        reveal_type(x)  # revealed: Unknown | A | B
     # This call would observe `x` as `A`.
     inner()
 
@@ -40,7 +37,7 @@ def outer(flag: bool) -> None:
     x = A()
 
     def inner() -> None:
-        reveal_type(x)  # revealed: Unknown | A | B | C
+        reveal_type(x)  # revealed: A | B | C
     inner()
 
     if flag:
@@ -62,7 +59,7 @@ def outer() -> None:
     x = A()
 
     def inner() -> None:
-        reveal_type(x)  # revealed: Unknown | A | C
+        reveal_type(x)  # revealed: A | C
     inner()
 
     if False:
@@ -76,7 +73,7 @@ def outer(flag: bool) -> None:
     x = A()
 
     def inner() -> None:
-        reveal_type(x)  # revealed: Unknown | A | C
+        reveal_type(x)  # revealed: A | C
     inner()
 
     if flag:
@@ -96,16 +93,10 @@ def outer(flag: bool) -> None:
         x = A()
 
         def inner() -> None:
-            reveal_type(x)  # revealed: Unknown | A
+            reveal_type(x)  # revealed: A
         inner()
 ```
 
-In the future, we may try to be smarter about which bindings must or must not be a visible to a
-given nested scope, depending where it is defined. In the above case, this shouldn't change the
-behavior -- `x` is defined before `inner` in the same branch, so should be considered
-definitely-bound for `inner`. But in other cases we may want to emit `possibly-unresolved-reference`
-in future:
-
 ```py
 def outer(flag: bool) -> None:
     if flag:
@@ -113,7 +104,7 @@ def outer(flag: bool) -> None:
 
     def inner() -> None:
         # TODO: Ideally, we would emit a possibly-unresolved-reference error here.
-        reveal_type(x)  # revealed: Unknown | A
+        reveal_type(x)  # revealed: A
     inner()
 ```
 
@@ -126,7 +117,7 @@ def outer() -> None:
     x = A()
 
     def inner() -> None:
-        reveal_type(x)  # revealed: Unknown | A
+        reveal_type(x)  # revealed: A
     inner()
 
     return
@@ -136,7 +127,7 @@ def outer(flag: bool) -> None:
     x = A()
 
     def inner() -> None:
-        reveal_type(x)  # revealed: Unknown | A | B
+        reveal_type(x)  # revealed: A | B
     if flag:
         x = B()
         inner()
@@ -161,7 +152,7 @@ def f0() -> None:
         def f2() -> None:
             def f3() -> None:
                 def f4() -> None:
-                    reveal_type(x)  # revealed: Unknown | A | B
+                    reveal_type(x)  # revealed: A | B
                 f4()
             f3()
         f2()
@@ -172,6 +163,29 @@ def f0() -> None:
     f1()
 ```
 
+## Narrowing
+
+In general, it is not safe to narrow the public type of a symbol using constraints introduced in an
+outer scope (because the symbol's value may have changed by the time the lazy scope is actually
+evaluated), but they can be applied if there is no reassignment of the symbol.
+
+```py
+class A: ...
+
+def outer(x: A | None):
+    if x is not None:
+        def inner() -> None:
+            reveal_type(x)  # revealed: A | None
+        inner()
+    x = None
+
+def outer(x: A | None):
+    if x is not None:
+        def inner() -> None:
+            reveal_type(x)  # revealed: A
+        inner()
+```
+
 ## At module level
 
 The behavior is the same if the outer scope is the global scope of a module:
@@ -232,32 +246,16 @@ def _():
 
 ## Limitations
 
-### Type narrowing
-
-We currently do not further analyze control flow, so we do not support cases where the inner scope
-is only executed in a branch where the type of `x` is narrowed:
-
-```py
-class A: ...
-
-def outer(x: A | None):
-    if x is not None:
-        def inner() -> None:
-            # TODO: should ideally be `A`
-            reveal_type(x)  # revealed: A | None
-        inner()
-```
-
 ### Shadowing
 
-Similarly, since we do not analyze control flow in the outer scope here, we assume that `inner()`
+Since we do not analyze control flow in the outer scope here, we assume that `inner()` could be
-could be called between the two assignments to `x`:
+called between the two assignments to `x`:
 
 ```py
 def outer() -> None:
     def inner() -> None:
-        # TODO: this should ideally be `Unknown | Literal[1]`, but no other type checker supports this either
+        # TODO: this should ideally be `Literal[1]`, but no other type checker supports this either
-        reveal_type(x)  # revealed: Unknown | None | Literal[1]
+        reveal_type(x)  # revealed: None | Literal[1]
     x = None
 
     # [additional code here]
@@ -279,8 +277,8 @@ def outer() -> None:
     x = 1
 
     def inner() -> None:
-        # TODO: this should be `Unknown | Literal[1]`. Mypy and pyright support this.
+        # TODO: this should be `Literal[1]`. Mypy and pyright support this.
-        reveal_type(x)  # revealed: Unknown | None | Literal[1]
+        reveal_type(x)  # revealed: None | Literal[1]
     inner()
 ```
 
@@ -314,8 +312,8 @@ def outer() -> None:
     set_x()
 
     def inner() -> None:
-        # TODO: this should ideally be `Unknown | None | Literal[1]`. Mypy and pyright support this.
+        # TODO: this should ideally be `None | Literal[1]`. Mypy and pyright support this.
-        reveal_type(x)  # revealed: Unknown | None
+        reveal_type(x)  # revealed: None
     inner()
 ```
 

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

@@ -299,7 +299,7 @@ def _():
     x = 1
 
     def f():
-        # revealed: Unknown | Literal[1, 2]
+        # revealed: Literal[1, 2]
         [reveal_type(x) for a in range(1)]
     x = 2
 ```
@@ -316,7 +316,7 @@ def _():
 
     class A:
         def f():
-            # revealed: Unknown | Literal[1, 2]
+            # revealed: Literal[1, 2]
             reveal_type(x)
 
     x = 2
@@ -333,7 +333,7 @@ def _():
 
     def f():
         def g():
-            # revealed: Unknown | Literal[1, 2]
+            # revealed: Literal[1, 2]
             reveal_type(x)
     x = 2
 ```
@@ -351,7 +351,7 @@ def _():
 
     class A:
         def f():
-            # revealed: Unknown | Literal[1, 2]
+            # revealed: Literal[1, 2]
             [reveal_type(x) for a in range(1)]
 
     x = 2

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

@@ -6,7 +6,7 @@
 def f():
     x = 1
     def g():
-        reveal_type(x)  # revealed: Unknown | Literal[1]
+        reveal_type(x)  # revealed: Literal[1]
 ```
 
 ## Two levels up
@@ -16,7 +16,7 @@ def f():
     x = 1
     def g():
         def h():
-            reveal_type(x)  # revealed: Unknown | Literal[1]
+            reveal_type(x)  # revealed: Literal[1]
 ```
 
 ## Skips class scope
@@ -28,7 +28,7 @@ def f():
     class C:
         x = 2
         def g():
-            reveal_type(x)  # revealed: Unknown | Literal[1]
+            reveal_type(x)  # revealed: Literal[1]
 ```
 
 ## Reads respect annotation-only declarations
@@ -104,12 +104,12 @@ def a():
 
             def d():
                 nonlocal x
-                reveal_type(x)  # revealed: Unknown | Literal[3, 2]
+                reveal_type(x)  # revealed: Literal[3, 2]
                 x = 4
                 reveal_type(x)  # revealed: Literal[4]
 
                 def e():
-                    reveal_type(x)  # revealed: Unknown | Literal[4, 3, 2]
+                    reveal_type(x)  # revealed: Literal[4, 3, 2]
 ```
 
 However, currently the union of types that we build is incomplete. We walk parent scopes, but not
@@ -127,7 +127,7 @@ def a():
             nonlocal x
             x = 3
         # TODO: This should include 2 and 3.
-        reveal_type(x)  # revealed: Unknown | Literal[1]
+        reveal_type(x)  # revealed: Literal[1]
 ```
 
 ## Local variable bindings "look ahead" to any assignment in the current scope
@@ -249,7 +249,7 @@ def f():
         x = 2
         def h():
             nonlocal x
-            reveal_type(x)  # revealed: Unknown | Literal[2]
+            reveal_type(x)  # revealed: Literal[2]
 ```
 
 ## `nonlocal` "chaining"
@@ -263,7 +263,7 @@ def f():
         nonlocal x
         def h():
             nonlocal x
-            reveal_type(x)  # revealed: Unknown | Literal[1]
+            reveal_type(x)  # revealed: Literal[1]
 ```
 
 And the `nonlocal` chain can skip over a scope that doesn't bind the variable:
@@ -277,7 +277,7 @@ def f1():
             # No binding; this scope gets skipped.
             def f4():
                 nonlocal x
-                reveal_type(x)  # revealed: Unknown | Literal[1]
+                reveal_type(x)  # revealed: Literal[1]
 ```
 
 But a `global` statement breaks the chain:
@@ -353,7 +353,7 @@ affected by `g`:
 def f():
     x = 1
     def g():
-        reveal_type(x)  # revealed: Unknown | Literal[1]
+        reveal_type(x)  # revealed: Literal[1]
     reveal_type(x)  # revealed: Literal[1]
 ```
 
@@ -365,9 +365,9 @@ def f():
     x = 1
     def g():
         nonlocal x
-        reveal_type(x)  # revealed: Unknown | Literal[1]
+        reveal_type(x)  # revealed: Literal[1]
         x += 1
-        reveal_type(x)  # revealed: Unknown | Literal[2]
+        reveal_type(x)  # revealed: Literal[2]
     # TODO: should be `Unknown | Literal[1]`
     reveal_type(x)  # revealed: Literal[1]
 ```
@@ -379,7 +379,7 @@ def f():
     x = 1
     def g():
         nonlocal x
-        reveal_type(x)  # revealed: Unknown | Literal[1]
+        reveal_type(x)  # revealed: Literal[1]
     # TODO: should be `Unknown | Literal[1]`
     reveal_type(x)  # revealed: Literal[1]
 ```

crates/ty_python_semantic/resources/mdtest/terminal_statements.md

@@ -806,11 +806,7 @@ def top_level_return(cond1: bool, cond2: bool):
     x = 1
 
     def g():
-        # TODO We could potentially eliminate `Unknown` from the union here,
+        reveal_type(x)  # revealed: Literal[1, 2, 3]
-        # because `x` resolves to an enclosing function-like scope and there
-        # are no nested `nonlocal` declarations of that symbol that might
-        # modify it.
-        reveal_type(x)  # revealed: Unknown | Literal[1, 2, 3]
     if cond1:
         if cond2:
             x = 2
@@ -822,7 +818,7 @@ def return_from_if(cond1: bool, cond2: bool):
     x = 1
 
     def g():
-        reveal_type(x)  # revealed: Unknown | Literal[1, 2, 3]
+        reveal_type(x)  # revealed: Literal[1, 2, 3]
     if cond1:
         if cond2:
             x = 2
@@ -834,7 +830,7 @@ def return_from_nested_if(cond1: bool, cond2: bool):
     x = 1
 
     def g():
-        reveal_type(x)  # revealed: Unknown | Literal[1, 2, 3]
+        reveal_type(x)  # revealed: Literal[1, 2, 3]
     if cond1:
         if cond2:
             x = 2

crates/ty_python_semantic/resources/mdtest/unreachable.md

@@ -250,7 +250,7 @@ def outer():
     x = 1
 
     def inner():
-        reveal_type(x)  # revealed: Unknown | Literal[1]
+        reveal_type(x)  # revealed: Literal[1]
     while True:
         pass
 ```

crates/ty_python_semantic/src/place.rs

@@ -768,12 +768,14 @@ fn place_by_id<'db>(
                 .expr
                 .is_name_and(|name| matches!(name, "__slots__" | "TYPE_CHECKING"));
 
-            if scope.file(db).is_stub(db) {
+            if scope.file(db).is_stub(db) || scope.scope(db).visibility().is_private() {
                 // We generally trust module-level undeclared places in stubs and do not union
                 // with `Unknown`. If we don't do this, simple aliases like `IOError = OSError` in
                 // stubs would result in `IOError` being a union of `OSError` and `Unknown`, which
                 // leads to all sorts of downstream problems. Similarly, type variables are often
                 // defined as `_T = TypeVar("_T")`, without being declared.
+                // Also, if the scope is private, such as a function scope,
+                // meaning that the place cannot be rewritten from elsewhere, we do not union with `Unknown`.
 
                 inferred.into()
             } else {

crates/ty_python_semantic/src/semantic_index.rs

@@ -24,7 +24,7 @@ use crate::semantic_index::place::{
     FileScopeId, NodeWithScopeKey, NodeWithScopeRef, PlaceExpr, PlaceTable, Scope, ScopeId,
     ScopeKind, ScopedPlaceId,
 };
-use crate::semantic_index::use_def::{EagerSnapshotKey, ScopedEagerSnapshotId, UseDefMap};
+use crate::semantic_index::use_def::{EnclosingSnapshotKey, ScopedEnclosingSnapshotId, UseDefMap};
 use crate::semantic_model::HasTrackedScope;
 use crate::util::get_size::untracked_arc_size;
 
@@ -188,7 +188,39 @@ pub(crate) fn global_scope(db: &dyn Db, file: File) -> ScopeId<'_> {
     FileScopeId::global().to_scope_id(db, file)
 }
 
-pub(crate) enum EagerSnapshotResult<'map, 'db> {
+#[derive(Debug, PartialEq, Eq, Clone, Copy, Hash, get_size2::GetSize)]
+pub(crate) enum ScopeVisibility {
+    /// The scope is private (e.g. function, type alias, comprehension scope).
+    Private,
+    /// The scope is public (e.g. module, class scope).
+    Public,
+}
+
+impl ScopeVisibility {
+    pub(crate) const fn is_public(self) -> bool {
+        matches!(self, ScopeVisibility::Public)
+    }
+
+    pub(crate) const fn is_private(self) -> bool {
+        matches!(self, ScopeVisibility::Private)
+    }
+}
+
+#[derive(Debug, PartialEq, Eq, Clone, Copy, Hash, get_size2::GetSize)]
+pub(crate) enum ScopeLaziness {
+    /// The scope is evaluated lazily (e.g. function, type alias scope).
+    Lazy,
+    /// The scope is evaluated eagerly (e.g. module, class, comprehension scope).
+    Eager,
+}
+
+impl ScopeLaziness {
+    pub(crate) const fn is_eager(self) -> bool {
+        matches!(self, ScopeLaziness::Eager)
+    }
+}
+
+pub(crate) enum EnclosingSnapshotResult<'map, 'db> {
     FoundConstraint(ScopedNarrowingConstraint),
     FoundBindings(BindingWithConstraintsIterator<'map, 'db>),
     NotFound,
@@ -234,8 +266,8 @@ pub(crate) struct SemanticIndex<'db> {
     /// Flags about the global scope (code usage impacting inference)
     has_future_annotations: bool,
 
-    /// Map of all of the eager snapshots that appear in this file.
+    /// Map of all of the enclosing snapshots that appear in this file.
-    eager_snapshots: FxHashMap<EagerSnapshotKey, ScopedEagerSnapshotId>,
+    enclosing_snapshots: FxHashMap<EnclosingSnapshotKey, ScopedEnclosingSnapshotId>,
 
     /// List of all semantic syntax errors in this file.
     semantic_syntax_errors: Vec<SemanticSyntaxError>,
@@ -484,36 +516,56 @@ impl<'db> SemanticIndex<'db> {
 
     /// Returns
     /// * `NoLongerInEagerContext` if the nested scope is no longer in an eager context
-    ///   (that is, not every scope that will be traversed is eager).
+    ///   (that is, not every scope that will be traversed is eager) and no lazy snapshots were found.
-    /// *  an iterator of bindings for a particular nested eager scope reference if the bindings exist.
+    /// *  an iterator of bindings for a particular nested scope reference if the bindings exist.
-    /// *  a narrowing constraint if there are no bindings, but there is a narrowing constraint for an outer scope symbol.
+    /// *  a narrowing constraint if there are no bindings, but there is a narrowing constraint for an enclosing scope place.
-    /// * `NotFound` if the narrowing constraint / bindings do not exist in the nested eager scope.
+    /// * `NotFound` if the narrowing constraint / bindings do not exist in the nested scope.
-    pub(crate) fn eager_snapshot(
+    pub(crate) fn enclosing_snapshot(
         &self,
         enclosing_scope: FileScopeId,
         expr: &PlaceExpr,
         nested_scope: FileScopeId,
-    ) -> EagerSnapshotResult<'_, 'db> {
+    ) -> EnclosingSnapshotResult<'_, 'db> {
         for (ancestor_scope_id, ancestor_scope) in self.ancestor_scopes(nested_scope) {
             if ancestor_scope_id == enclosing_scope {
                 break;
             }
             if !ancestor_scope.is_eager() {
-                return EagerSnapshotResult::NoLongerInEagerContext;
+                if expr.is_name() {
+                    if let Some(place_id) =
+                        self.place_tables[enclosing_scope].place_id_by_expr(expr)
+                    {
+                        let key = EnclosingSnapshotKey {
+                            enclosing_scope,
+                            enclosing_place: place_id,
+                            nested_scope,
+                            nested_laziness: ScopeLaziness::Lazy,
+                        };
+                        if let Some(id) = self.enclosing_snapshots.get(&key) {
+                            return self.use_def_maps[enclosing_scope]
+                                .inner
+                                .enclosing_snapshot(*id, key.nested_laziness);
+                        }
+                    }
+                }
+                return EnclosingSnapshotResult::NoLongerInEagerContext;
             }
         }
         let Some(place_id) = self.place_tables[enclosing_scope].place_id_by_expr(expr) else {
-            return EagerSnapshotResult::NotFound;
+            return EnclosingSnapshotResult::NotFound;
         };
-        let key = EagerSnapshotKey {
+        let key = EnclosingSnapshotKey {
             enclosing_scope,
             enclosing_place: place_id,
             nested_scope,
+            nested_laziness: ScopeLaziness::Eager,
         };
-        let Some(id) = self.eager_snapshots.get(&key) else {
+        let Some(id) = self.enclosing_snapshots.get(&key) else {
-            return EagerSnapshotResult::NotFound;
+            return EnclosingSnapshotResult::NotFound;
         };
-        self.use_def_maps[enclosing_scope].inner.eager_snapshot(*id)
+        self.use_def_maps[enclosing_scope]
+            .inner
+            .enclosing_snapshot(*id, key.nested_laziness)
     }
 
     pub(crate) fn semantic_syntax_errors(&self) -> &[SemanticSyntaxError] {

crates/ty_python_semantic/src/semantic_index/builder.rs

@@ -43,9 +43,9 @@ use crate::semantic_index::reachability_constraints::{
     ReachabilityConstraintsBuilder, ScopedReachabilityConstraintId,
 };
 use crate::semantic_index::use_def::{
-    EagerSnapshotKey, FlowSnapshot, ScopedEagerSnapshotId, UseDefMapBuilder,
+    EnclosingSnapshotKey, FlowSnapshot, ScopedEnclosingSnapshotId, UseDefMapBuilder,
 };
-use crate::semantic_index::{ArcUseDefMap, ExpressionsScopeMap, SemanticIndex};
+use crate::semantic_index::{ArcUseDefMap, ExpressionsScopeMap, ScopeLaziness, SemanticIndex};
 use crate::semantic_model::HasTrackedScope;
 use crate::unpack::{Unpack, UnpackKind, UnpackPosition, UnpackValue};
 use crate::{Db, Program};
@@ -113,7 +113,7 @@ pub(super) struct SemanticIndexBuilder<'db, 'ast> {
     ///
     /// [generator functions]: https://docs.python.org/3/glossary.html#term-generator
     generator_functions: FxHashSet<FileScopeId>,
-    eager_snapshots: FxHashMap<EagerSnapshotKey, ScopedEagerSnapshotId>,
+    enclosing_snapshots: FxHashMap<EnclosingSnapshotKey, ScopedEnclosingSnapshotId>,
     /// Errors collected by the `semantic_checker`.
     semantic_syntax_errors: RefCell<Vec<SemanticSyntaxError>>,
 }
@@ -148,7 +148,7 @@ impl<'db, 'ast> SemanticIndexBuilder<'db, 'ast> {
             imported_modules: FxHashSet::default(),
             generator_functions: FxHashSet::default(),
 
-            eager_snapshots: FxHashMap::default(),
+            enclosing_snapshots: FxHashMap::default(),
 
             python_version: Program::get(db).python_version(db),
             source_text: OnceCell::new(),
@@ -276,25 +276,8 @@ impl<'db, 'ast> SemanticIndexBuilder<'db, 'ast> {
         });
     }
 
-    fn pop_scope(&mut self) -> FileScopeId {
+    // Records snapshots of the place states visible from the current eager scope.
-        self.try_node_context_stack_manager.exit_scope();
+    fn record_eager_snapshots(&mut self, popped_scope_id: FileScopeId) {
-
-        let ScopeInfo {
-            file_scope_id: popped_scope_id,
-            ..
-        } = self
-            .scope_stack
-            .pop()
-            .expect("Root scope should be present");
-
-        let children_end = self.scopes.next_index();
-        let popped_scope = &mut self.scopes[popped_scope_id];
-        popped_scope.extend_descendants(children_end);
-
-        if !popped_scope.is_eager() {
-            return 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.
@@ -327,27 +310,163 @@ impl<'db, 'ast> SemanticIndexBuilder<'db, 'ast> {
 
                 // Snapshot the state of this place that are visible at this point in this
                 // enclosing scope.
-                let key = EagerSnapshotKey {
+                let key = EnclosingSnapshotKey {
                     enclosing_scope: enclosing_scope_id,
                     enclosing_place: enclosing_place_id,
                     nested_scope: popped_scope_id,
+                    nested_laziness: ScopeLaziness::Eager,
                 };
-                let eager_snapshot = self.use_def_maps[enclosing_scope_id].snapshot_eager_state(
+                let eager_snapshot = self.use_def_maps[enclosing_scope_id].snapshot_outer_state(
                     enclosing_place_id,
                     enclosing_scope_kind,
                     enclosing_place,
                 );
-                self.eager_snapshots.insert(key, eager_snapshot);
+                self.enclosing_snapshots.insert(key, eager_snapshot);
             }
 
             // Lazy scopes are "sticky": once we see a lazy scope we stop doing lookups
             // eagerly, even if we would encounter another eager enclosing scope later on.
-            // Also, narrowing constraints outside a lazy scope are not applicable.
-            // TODO: If the place has never been rewritten, they are applicable.
             if !enclosing_scope_kind.is_eager() {
                 break;
             }
         }
+    }
+
+    fn bound_scope(
+        &self,
+        enclosing_scope: FileScopeId,
+        place_expr: &PlaceExpr,
+    ) -> Option<FileScopeId> {
+        self.scope_stack
+            .iter()
+            .rev()
+            .skip_while(|scope| scope.file_scope_id != enclosing_scope)
+            .find_map(|scope_info| {
+                let scope_id = scope_info.file_scope_id;
+                let place_table = &self.place_tables[scope_id];
+                let place_id = place_table.place_id_by_expr(place_expr)?;
+                if place_table.place_expr(place_id).is_bound() {
+                    Some(scope_id)
+                } else {
+                    None
+                }
+            })
+    }
+
+    // Records snapshots of the place states visible from the current lazy scope.
+    fn record_lazy_snapshots(&mut self, popped_scope_id: FileScopeId) {
+        for enclosing_scope_info in self.scope_stack.iter().rev() {
+            let enclosing_scope_id = enclosing_scope_info.file_scope_id;
+            let enclosing_scope_kind = self.scopes[enclosing_scope_id].kind();
+            let enclosing_place_table = &self.place_tables[enclosing_scope_id];
+
+            for nested_place in self.place_tables[popped_scope_id].places() {
+                // We don't record lazy snapshots of attributes or subscripts, because these are difficult to track as they modify.
+                // For the same reason, symbols declared as nonlocal or global are not recorded.
+                // Also, if the enclosing scope allows its members to be modified from elsewhere, the snapshot will not be recorded.
+                if !nested_place.is_name()
+                    || self.scopes[enclosing_scope_id].visibility().is_public()
+                {
+                    continue;
+                }
+
+                // Skip this place if this enclosing scope doesn't contain any bindings for it.
+                // Note that even if this place is bound in the popped scope,
+                // it may refer to the enclosing scope bindings
+                // so we also need to snapshot the bindings of the enclosing scope.
+
+                let Some(enclosing_place_id) =
+                    enclosing_place_table.place_id_by_expr(&nested_place.expr)
+                else {
+                    continue;
+                };
+                let enclosing_place = enclosing_place_table.place_expr(enclosing_place_id);
+                if !enclosing_place.is_bound() {
+                    // If the bound scope of a place can be modified from elsewhere, the snapshot will not be recorded.
+                    if self
+                        .bound_scope(enclosing_scope_id, &nested_place.expr)
+                        .is_none_or(|scope| self.scopes[scope].visibility().is_public())
+                    {
+                        continue;
+                    }
+                }
+
+                // Snapshot the state of this place that are visible at this point in this
+                // enclosing scope (this may later be invalidated and swept away).
+                let key = EnclosingSnapshotKey {
+                    enclosing_scope: enclosing_scope_id,
+                    enclosing_place: enclosing_place_id,
+                    nested_scope: popped_scope_id,
+                    nested_laziness: ScopeLaziness::Lazy,
+                };
+                let lazy_snapshot = self.use_def_maps[enclosing_scope_id].snapshot_outer_state(
+                    enclosing_place_id,
+                    enclosing_scope_kind,
+                    enclosing_place,
+                );
+                self.enclosing_snapshots.insert(key, lazy_snapshot);
+            }
+        }
+    }
+
+    /// Any lazy snapshots of places that have been reassigned or modified are no longer valid, so delete them.
+    fn sweep_lazy_snapshots(&mut self, popped_scope_id: FileScopeId) {
+        self.enclosing_snapshots.retain(|key, _| {
+            let place_table = &self.place_tables[key.enclosing_scope];
+            key.nested_laziness.is_eager()
+                || key.enclosing_scope != popped_scope_id
+                || !place_table.is_place_reassigned(key.enclosing_place)
+        });
+    }
+
+    fn sweep_nonlocal_lazy_snapshots(&mut self) {
+        self.enclosing_snapshots.retain(|key, _| {
+            let place_table = &self.place_tables[key.enclosing_scope];
+
+            let is_place_bound_and_nonlocal = || -> bool {
+                let place_expr = place_table.place_expr(key.enclosing_place);
+                self.scopes
+                    .iter_enumerated()
+                    .skip_while(|(scope_id, _)| *scope_id != key.enclosing_scope)
+                    .any(|(scope_id, _)| {
+                        let other_scope_place_table = &self.place_tables[scope_id];
+                        let Some(place_id) =
+                            other_scope_place_table.place_id_by_expr(&place_expr.expr)
+                        else {
+                            return false;
+                        };
+                        let place = other_scope_place_table.place_expr(place_id);
+                        place.is_marked_nonlocal() && place.is_bound()
+                    })
+            };
+
+            key.nested_laziness.is_eager() || !is_place_bound_and_nonlocal()
+        });
+    }
+
+    fn pop_scope(&mut self) -> FileScopeId {
+        self.try_node_context_stack_manager.exit_scope();
+
+        let ScopeInfo {
+            file_scope_id: popped_scope_id,
+            ..
+        } = self
+            .scope_stack
+            .pop()
+            .expect("Root scope should be present");
+
+        self.sweep_lazy_snapshots(popped_scope_id);
+
+        let children_end = self.scopes.next_index();
+
+        let popped_scope = &mut self.scopes[popped_scope_id];
+        popped_scope.extend_descendants(children_end);
+
+        if popped_scope.is_eager() {
+            self.record_eager_snapshots(popped_scope_id);
+        } else {
+            self.record_lazy_snapshots(popped_scope_id);
+        }
 
         popped_scope_id
     }
@@ -1037,6 +1156,7 @@ impl<'db, 'ast> SemanticIndexBuilder<'db, 'ast> {
 
         // Pop the root scope
         self.pop_scope();
+        self.sweep_nonlocal_lazy_snapshots();
         assert!(self.scope_stack.is_empty());
 
         assert_eq!(&self.current_assignments, &[]);
@@ -1076,7 +1196,7 @@ impl<'db, 'ast> SemanticIndexBuilder<'db, 'ast> {
         self.scope_ids_by_scope.shrink_to_fit();
         self.scopes_by_node.shrink_to_fit();
         self.generator_functions.shrink_to_fit();
-        self.eager_snapshots.shrink_to_fit();
+        self.enclosing_snapshots.shrink_to_fit();
 
         SemanticIndex {
             place_tables,
@@ -1090,7 +1210,7 @@ impl<'db, 'ast> SemanticIndexBuilder<'db, 'ast> {
             use_def_maps,
             imported_modules: Arc::new(self.imported_modules),
             has_future_annotations: self.has_future_annotations,
-            eager_snapshots: self.eager_snapshots,
+            enclosing_snapshots: self.enclosing_snapshots,
             semantic_syntax_errors: self.semantic_syntax_errors.into_inner(),
             generator_functions: self.generator_functions,
         }

crates/ty_python_semantic/src/semantic_index/narrowing_constraints.rs

@@ -43,7 +43,7 @@ pub(crate) type ScopedNarrowingConstraint = List<ScopedNarrowingConstraintPredic
 #[derive(Clone, Copy, Debug, Eq, PartialEq)]
 pub(crate) enum ConstraintKey {
     NarrowingConstraint(ScopedNarrowingConstraint),
-    EagerNestedScope(FileScopeId),
+    NestedScope(FileScopeId),
     UseId(ScopedUseId),
 }
 

crates/ty_python_semantic/src/semantic_index/place.rs

@@ -16,7 +16,9 @@ use crate::Db;
 use crate::ast_node_ref::AstNodeRef;
 use crate::node_key::NodeKey;
 use crate::semantic_index::reachability_constraints::ScopedReachabilityConstraintId;
-use crate::semantic_index::{PlaceSet, SemanticIndex, semantic_index};
+use crate::semantic_index::{
+    PlaceSet, ScopeLaziness, ScopeVisibility, SemanticIndex, semantic_index,
+};
 
 #[derive(Debug, Clone, PartialEq, Eq, Hash, get_size2::GetSize)]
 pub(crate) enum PlaceExprSubSegment {
@@ -349,6 +351,10 @@ impl PlaceExprWithFlags {
         self.flags.contains(PlaceFlags::MARKED_NONLOCAL)
     }
 
+    pub fn is_reassigned(&self) -> bool {
+        self.flags.contains(PlaceFlags::IS_REASSIGNED)
+    }
+
     pub(crate) fn as_name(&self) -> Option<&Name> {
         self.expr.as_name()
     }
@@ -419,6 +425,7 @@ bitflags! {
         const MARKED_GLOBAL         = 1 << 3;
         const MARKED_NONLOCAL       = 1 << 4;
         const IS_INSTANCE_ATTRIBUTE = 1 << 5;
+        const IS_REASSIGNED         = 1 << 6;
     }
 }
 
@@ -579,6 +586,10 @@ impl Scope {
         self.node().scope_kind()
     }
 
+    pub(crate) fn visibility(&self) -> ScopeVisibility {
+        self.kind().visibility()
+    }
+
     pub fn descendants(&self) -> Range<FileScopeId> {
         self.descendants.clone()
     }
@@ -613,12 +624,27 @@ pub enum ScopeKind {
 
 impl ScopeKind {
     pub(crate) const fn is_eager(self) -> bool {
+        self.laziness().is_eager()
+    }
+
+    pub(crate) const fn laziness(self) -> ScopeLaziness {
         match self {
-            ScopeKind::Module | ScopeKind::Class | ScopeKind::Comprehension => true,
+            ScopeKind::Module | ScopeKind::Class | ScopeKind::Comprehension => ScopeLaziness::Eager,
             ScopeKind::Annotation
             | ScopeKind::Function
             | ScopeKind::Lambda
-            | ScopeKind::TypeAlias => false,
+            | ScopeKind::TypeAlias => ScopeLaziness::Lazy,
+        }
+    }
+
+    pub(crate) const fn visibility(self) -> ScopeVisibility {
+        match self {
+            ScopeKind::Module | ScopeKind::Class => ScopeVisibility::Public,
+            ScopeKind::Annotation
+            | ScopeKind::TypeAlias
+            | ScopeKind::Function
+            | ScopeKind::Lambda
+            | ScopeKind::Comprehension => ScopeVisibility::Private,
         }
     }
 
@@ -842,6 +868,9 @@ impl PlaceTableBuilder {
     }
 
     pub(super) fn mark_place_bound(&mut self, id: ScopedPlaceId) {
+        if self.table.places[id].is_bound() {
+            self.table.places[id].insert_flags(PlaceFlags::IS_REASSIGNED);
+        }
         self.table.places[id].insert_flags(PlaceFlags::IS_BOUND);
     }
 
@@ -873,6 +902,10 @@ impl PlaceTableBuilder {
         self.table.place_expr(place_id)
     }
 
+    pub(super) fn is_place_reassigned(&self, place_id: ScopedPlaceId) -> bool {
+        self.table.places[place_id].is_reassigned()
+    }
+
     /// Returns the place IDs associated with the place (e.g. `x.y`, `x.y.z`, `x.y.z[0]` for `x`).
     pub(super) fn associated_place_ids(
         &self,

crates/ty_python_semantic/src/semantic_index/use_def.rs

@@ -244,7 +244,7 @@ use ruff_index::{IndexVec, newtype_index};
 use rustc_hash::FxHashMap;
 
 use self::place_state::{
-    Bindings, Declarations, EagerSnapshot, LiveBindingsIterator, LiveDeclaration,
+    Bindings, Declarations, EnclosingSnapshot, LiveBindingsIterator, LiveDeclaration,
     LiveDeclarationsIterator, PlaceState, ScopedDefinitionId,
 };
 use crate::node_key::NodeKey;
@@ -264,7 +264,7 @@ use crate::semantic_index::reachability_constraints::{
     ReachabilityConstraints, ReachabilityConstraintsBuilder, ScopedReachabilityConstraintId,
 };
 use crate::semantic_index::use_def::place_state::PreviousDefinitions;
-use crate::semantic_index::{EagerSnapshotResult, SemanticIndex};
+use crate::semantic_index::{EnclosingSnapshotResult, ScopeLaziness, SemanticIndex};
 use crate::types::{IntersectionBuilder, Truthiness, Type, infer_narrowing_constraint};
 
 mod place_state;
@@ -318,8 +318,8 @@ pub(crate) struct UseDefMap<'db> {
     reachable_definitions: IndexVec<ScopedPlaceId, ReachableDefinitions>,
 
     /// Snapshot of bindings in this scope that can be used to resolve a reference in a nested
-    /// eager scope.
+    /// scope.
-    eager_snapshots: EagerSnapshots,
+    enclosing_snapshots: EnclosingSnapshots,
 
     /// Whether or not the end of the scope is reachable.
     ///
@@ -371,9 +371,9 @@ impl<'db> UseDefMap<'db> {
                     constraint_ids: self.narrowing_constraints.iter_predicates(constraint),
                 })
             }
-            ConstraintKey::EagerNestedScope(nested_scope) => {
+            ConstraintKey::NestedScope(nested_scope) => {
-                let EagerSnapshotResult::FoundBindings(bindings) =
+                let EnclosingSnapshotResult::FoundBindings(bindings) =
-                    index.eager_snapshot(enclosing_scope, expr, nested_scope)
+                    index.enclosing_snapshot(enclosing_scope, expr, nested_scope)
                 else {
                     unreachable!(
                         "The result of `SemanticIndex::eager_snapshot` must be `FoundBindings`"
@@ -436,18 +436,25 @@ impl<'db> UseDefMap<'db> {
         )
     }
 
-    pub(crate) fn eager_snapshot(
+    pub(crate) fn enclosing_snapshot(
         &self,
-        eager_bindings: ScopedEagerSnapshotId,
+        snapshot_id: ScopedEnclosingSnapshotId,
-    ) -> EagerSnapshotResult<'_, 'db> {
+        nested_laziness: ScopeLaziness,
-        match self.eager_snapshots.get(eager_bindings) {
+    ) -> EnclosingSnapshotResult<'_, 'db> {
-            Some(EagerSnapshot::Constraint(constraint)) => {
+        let boundness_analysis = if nested_laziness.is_eager() {
-                EagerSnapshotResult::FoundConstraint(*constraint)
+            BoundnessAnalysis::BasedOnUnboundVisibility
+        } else {
+            // TODO: We haven't implemented proper boundness analysis for nonlocal symbols, so we assume the boundness is bound for now.
+            BoundnessAnalysis::AssumeBound
+        };
+        match self.enclosing_snapshots.get(snapshot_id) {
+            Some(EnclosingSnapshot::Constraint(constraint)) => {
+                EnclosingSnapshotResult::FoundConstraint(*constraint)
             }
-            Some(EagerSnapshot::Bindings(bindings)) => EagerSnapshotResult::FoundBindings(
+            Some(EnclosingSnapshot::Bindings(bindings)) => EnclosingSnapshotResult::FoundBindings(
-                self.bindings_iterator(bindings, BoundnessAnalysis::BasedOnUnboundVisibility),
+                self.bindings_iterator(bindings, boundness_analysis),
             ),
-            None => EagerSnapshotResult::NotFound,
+            None => EnclosingSnapshotResult::NotFound,
         }
     }
 
@@ -566,30 +573,37 @@ impl<'db> UseDefMap<'db> {
     }
 }
 
-/// Uniquely identifies a snapshot of a place state that can be used to resolve a reference in a
+/// Uniquely identifies a snapshot of an enclosing scope place state that can be used to resolve a reference in a
-/// nested eager scope.
+/// nested scope.
 ///
 /// An eager scope has its entire body executed immediately at the location where it is defined.
 /// For any free references in the nested scope, we use the bindings that are visible at the point
 /// where the nested scope is defined, instead of using the public type of the place.
 ///
-/// There is a unique ID for each distinct [`EagerSnapshotKey`] in the file.
+/// There is a unique ID for each distinct [`EnclosingSnapshotKey`] in the file.
 #[newtype_index]
 #[derive(get_size2::GetSize)]
-pub(crate) struct ScopedEagerSnapshotId;
+pub(crate) struct ScopedEnclosingSnapshotId;
 
 #[derive(Clone, Copy, Debug, Eq, Hash, PartialEq, get_size2::GetSize)]
-pub(crate) struct EagerSnapshotKey {
+pub(crate) struct EnclosingSnapshotKey {
     /// The enclosing scope containing the bindings
     pub(crate) enclosing_scope: FileScopeId,
     /// The referenced place (in the enclosing scope)
     pub(crate) enclosing_place: ScopedPlaceId,
-    /// The nested eager scope containing the reference
+    /// The nested scope containing the reference
     pub(crate) nested_scope: FileScopeId,
+    /// Laziness of the nested scope
+    pub(crate) nested_laziness: ScopeLaziness,
 }
 
-/// A snapshot of place states that can be used to resolve a reference in a nested eager scope.
+/// A snapshot of enclosing scope place states that can be used to resolve a reference in a nested scope.
-type EagerSnapshots = IndexVec<ScopedEagerSnapshotId, EagerSnapshot>;
+/// Normally, if the current scope is lazily evaluated,
+/// we do not snapshot the place states from the enclosing scope,
+/// and infer the type of the place from its reachable definitions
+/// (and any narrowing constraints introduced in the enclosing scope do not apply to the current scope).
+/// The exception is if the symbol has never been reassigned, in which case it is snapshotted.
+type EnclosingSnapshots = IndexVec<ScopedEnclosingSnapshotId, EnclosingSnapshot>;
 
 #[derive(Debug)]
 pub(crate) struct BindingWithConstraintsIterator<'map, 'db> {
@@ -757,8 +771,8 @@ pub(super) struct UseDefMapBuilder<'db> {
     reachable_definitions: IndexVec<ScopedPlaceId, ReachableDefinitions>,
 
     /// Snapshots of place states in this scope that can be used to resolve a reference in a
-    /// nested eager scope.
+    /// nested scope.
-    eager_snapshots: EagerSnapshots,
+    enclosing_snapshots: EnclosingSnapshots,
 
     /// Is this a class scope?
     is_class_scope: bool,
@@ -778,10 +792,11 @@ impl<'db> UseDefMapBuilder<'db> {
             bindings_by_definition: FxHashMap::default(),
             place_states: IndexVec::new(),
             reachable_definitions: IndexVec::new(),
-            eager_snapshots: EagerSnapshots::default(),
+            enclosing_snapshots: EnclosingSnapshots::default(),
             is_class_scope,
         }
     }
+
     pub(super) fn mark_unreachable(&mut self) {
         self.reachability = ScopedReachabilityConstraintId::ALWAYS_FALSE;
 
@@ -1022,23 +1037,23 @@ impl<'db> UseDefMapBuilder<'db> {
         self.node_reachability.insert(node_key, self.reachability);
     }
 
-    pub(super) fn snapshot_eager_state(
+    pub(super) fn snapshot_outer_state(
         &mut self,
         enclosing_place: ScopedPlaceId,
         scope: ScopeKind,
         enclosing_place_expr: &PlaceExprWithFlags,
-    ) -> ScopedEagerSnapshotId {
+    ) -> ScopedEnclosingSnapshotId {
         // 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_expr.is_name()) || !enclosing_place_expr.is_bound()
         {
-            self.eager_snapshots.push(EagerSnapshot::Constraint(
+            self.enclosing_snapshots.push(EnclosingSnapshot::Constraint(
                 self.place_states[enclosing_place]
                     .bindings()
                     .unbound_narrowing_constraint(),
             ))
         } else {
-            self.eager_snapshots.push(EagerSnapshot::Bindings(
+            self.enclosing_snapshots.push(EnclosingSnapshot::Bindings(
                 self.place_states[enclosing_place].bindings().clone(),
             ))
         }
@@ -1144,7 +1159,7 @@ impl<'db> UseDefMapBuilder<'db> {
         for bindings in self.bindings_by_definition.values_mut() {
             bindings.finish(&mut self.reachability_constraints);
         }
-        for eager_snapshot in &mut self.eager_snapshots {
+        for eager_snapshot in &mut self.enclosing_snapshots {
             eager_snapshot.finish(&mut self.reachability_constraints);
         }
         self.reachability_constraints.mark_used(self.reachability);
@@ -1160,7 +1175,7 @@ impl<'db> UseDefMapBuilder<'db> {
         self.node_reachability.shrink_to_fit();
         self.declarations_by_binding.shrink_to_fit();
         self.bindings_by_definition.shrink_to_fit();
-        self.eager_snapshots.shrink_to_fit();
+        self.enclosing_snapshots.shrink_to_fit();
 
         UseDefMap {
             all_definitions: self.all_definitions,
@@ -1173,7 +1188,7 @@ impl<'db> UseDefMapBuilder<'db> {
             reachable_definitions: self.reachable_definitions,
             declarations_by_binding: self.declarations_by_binding,
             bindings_by_definition: self.bindings_by_definition,
-            eager_snapshots: self.eager_snapshots,
+            enclosing_snapshots: self.enclosing_snapshots,
             end_of_scope_reachability: self.reachability,
         }
     }

crates/ty_python_semantic/src/semantic_index/use_def/place_state.rs

@@ -181,22 +181,22 @@ impl Declarations {
     }
 }
 
-/// A snapshot of a place state that can be used to resolve a reference in a nested eager scope.
+/// A snapshot of a place state that can be used to resolve a reference in a nested scope.
-/// If there are bindings in a (non-class) scope , they are stored in `Bindings`.
+/// If there are bindings in a (non-class) scope, they are stored in `Bindings`.
 /// Even if it's a class scope (class variables are not visible to nested scopes) or there are no
 /// bindings, the current narrowing constraint is necessary for narrowing, so it's stored in
 /// `Constraint`.
 #[derive(Clone, Debug, PartialEq, Eq, salsa::Update, get_size2::GetSize)]
-pub(super) enum EagerSnapshot {
+pub(super) enum EnclosingSnapshot {
     Constraint(ScopedNarrowingConstraint),
     Bindings(Bindings),
 }
 
-impl EagerSnapshot {
+impl EnclosingSnapshot {
     pub(super) fn finish(&mut self, reachability_constraints: &mut ReachabilityConstraintsBuilder) {
         match self {
-            EagerSnapshot::Constraint(_) => {}
+            Self::Constraint(_) => {}
-            EagerSnapshot::Bindings(bindings) => {
+            Self::Bindings(bindings) => {
                 bindings.finish(reachability_constraints);
             }
         }

crates/ty_python_semantic/src/types/infer.rs

@@ -85,7 +85,7 @@ use crate::semantic_index::place::{
     FileScopeId, NodeWithScopeKind, NodeWithScopeRef, PlaceExpr, ScopeId, ScopeKind, ScopedPlaceId,
 };
 use crate::semantic_index::{
-    ApplicableConstraints, EagerSnapshotResult, SemanticIndex, place_table, semantic_index,
+    ApplicableConstraints, EnclosingSnapshotResult, SemanticIndex, place_table, semantic_index,
 };
 use crate::types::call::{Binding, Bindings, CallArguments, CallError};
 use crate::types::class::{CodeGeneratorKind, DataclassField, MetaclassErrorKind, SliceLiteral};
@@ -6461,17 +6461,18 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> {
                 // enclosing scopes that actually contain bindings that we should use when
                 // resolving the reference.)
                 if !self.is_deferred() {
-                    match self
+                    match self.index.enclosing_snapshot(
-                        .index
+                        enclosing_scope_file_id,
-                        .eager_snapshot(enclosing_scope_file_id, expr, file_scope_id)
+                        expr,
-                    {
+                        file_scope_id,
-                        EagerSnapshotResult::FoundConstraint(constraint) => {
+                    ) {
+                        EnclosingSnapshotResult::FoundConstraint(constraint) => {
                             constraint_keys.push((
                                 enclosing_scope_file_id,
                                 ConstraintKey::NarrowingConstraint(constraint),
                             ));
                         }
-                        EagerSnapshotResult::FoundBindings(bindings) => {
+                        EnclosingSnapshotResult::FoundBindings(bindings) => {
                             if expr.is_name()
                                 && !enclosing_scope_id.is_function_like(db)
                                 && !is_immediately_enclosing_scope
@@ -6487,13 +6488,13 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> {
                             });
                             constraint_keys.push((
                                 enclosing_scope_file_id,
-                                ConstraintKey::EagerNestedScope(file_scope_id),
+                                ConstraintKey::NestedScope(file_scope_id),
                             ));
                             return place.into();
                         }
                         // There are no visible bindings / constraint here.
                         // Don't fall back to non-eager place resolution.
-                        EagerSnapshotResult::NotFound => {
+                        EnclosingSnapshotResult::NotFound => {
                             let enclosing_place_table =
                                 self.index.place_table(enclosing_scope_file_id);
                             for enclosing_root_place in enclosing_place_table.root_place_exprs(expr)
@@ -6513,7 +6514,7 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> {
                             }
                             continue;
                         }
-                        EagerSnapshotResult::NoLongerInEagerContext => {}
+                        EnclosingSnapshotResult::NoLongerInEagerContext => {}
                     }
                 }
 
@@ -6583,17 +6584,18 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> {
                     }
 
                     if !self.is_deferred() {
-                        match self
+                        match self.index.enclosing_snapshot(
-                            .index
+                            FileScopeId::global(),
-                            .eager_snapshot(FileScopeId::global(), expr, file_scope_id)
+                            expr,
-                        {
+                            file_scope_id,
-                            EagerSnapshotResult::FoundConstraint(constraint) => {
+                        ) {
+                            EnclosingSnapshotResult::FoundConstraint(constraint) => {
                                 constraint_keys.push((
                                     FileScopeId::global(),
                                     ConstraintKey::NarrowingConstraint(constraint),
                                 ));
                             }
-                            EagerSnapshotResult::FoundBindings(bindings) => {
+                            EnclosingSnapshotResult::FoundBindings(bindings) => {
                                 let place = place_from_bindings(db, bindings).map_type(|ty| {
                                     self.narrow_place_with_applicable_constraints(
                                         expr,
@@ -6603,15 +6605,15 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> {
                                 });
                                 constraint_keys.push((
                                     FileScopeId::global(),
-                                    ConstraintKey::EagerNestedScope(file_scope_id),
+                                    ConstraintKey::NestedScope(file_scope_id),
                                 ));
                                 return place.into();
                             }
                             // There are no visible bindings / constraint here.
-                            EagerSnapshotResult::NotFound => {
+                            EnclosingSnapshotResult::NotFound => {
                                 return Place::Unbound.into();
                             }
-                            EagerSnapshotResult::NoLongerInEagerContext => {}
+                            EnclosingSnapshotResult::NoLongerInEagerContext => {}
                         }
                     }