[ty] more precise lazy scope place lookup

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

@@ -248,6 +248,13 @@ def f(x: str | None):
                 reveal_type(x)  # revealed: str | None
     x = None
 
+def f(x: str | None):
+    def _(x: str | None):
+        if x is not None:
+            def closure():
+                reveal_type(x)  # revealed: str
+    x = None
+
 def f(x: str | None):
     class C:
         def _():
@@ -303,13 +310,17 @@ no longer valid in the inner lazy scope.
 def f(l: list[str | None]):
     if l[0] is not None:
         def _():
-            reveal_type(l[0])  # revealed: str | None
+            # TODO: should be `str | None`
+            reveal_type(l[0])  # revealed: str | None | @Todo(list literal element type)
+        # TODO: should be of type `list[None]`
         l = [None]
 
 def f(l: list[str | None]):
     l[0] = "a"
     def _():
-        reveal_type(l[0])  # revealed: str | None
+        # TODO: should be `str | None`
+        reveal_type(l[0])  # revealed: str | None | @Todo(list literal element type)
+    # TODO: should be of type `list[None]`
     l = [None]
 
 def f(l: list[str | None]):

crates/ty_python_semantic/resources/mdtest/public_types.md

@@ -83,6 +83,18 @@ def outer(flag: bool) -> None:
 
     x = C()
     inner()
+
+def outer(flag: bool) -> None:
+    if flag:
+        x = A()
+    else:
+        x = B()  # this binding of `x` is invisible to `inner`
+        return
+
+    def inner() -> None:
+        reveal_type(x)  # revealed: A | C
+    x = C()
+    inner()
 ```
 
 If a symbol is only conditionally bound, we do not raise any errors:
@@ -186,6 +198,59 @@ def outer(x: A | None):
         inner()
 ```
 
+"Reassignment" here refers to a thing that happens after the closure is defined that can actually
+change the inferred type of a captured symbol. Something done before the closure definition is more
+of a shadowing, and doesn't actually invalidate narrowing.
+
+```py
+def outer() -> None:
+    x = None
+
+    def inner() -> None:
+        # In this scope, `x` may refer to `x = None` or `x = 1`.
+        reveal_type(x)  # revealed: None | Literal[1]
+    inner()
+
+    x = 1
+
+    inner()
+
+    def inner2() -> None:
+        # In this scope, `x = None` appears as being shadowed by `x = 1`.
+        reveal_type(x)  # revealed: Literal[1]
+    inner2()
+
+def outer() -> None:
+    x = None
+
+    x = 1
+
+    def inner() -> None:
+        reveal_type(x)  # revealed: Literal[1, 2]
+    inner()
+
+    x = 2
+
+def outer(x: A | None):
+    if x is None:
+        x = A()
+
+    reveal_type(x)  # revealed: A
+
+    def inner() -> None:
+        reveal_type(x)  # revealed: A
+    inner()
+
+def outer(x: A | None):
+    x = x or A()
+
+    reveal_type(x)  # revealed: A
+
+    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:
@@ -265,37 +330,17 @@ def outer() -> None:
     inner()
 ```
 
-This is currently even true if the `inner` function is only defined after the second assignment to
-`x`:
-
-```py
-def outer() -> None:
-    x = None
-
-    # [additional code here]
-
-    x = 1
-
-    def inner() -> None:
-        # TODO: this should be `Literal[1]`. Mypy and pyright support this.
-        reveal_type(x)  # revealed: None | Literal[1]
-    inner()
-```
-
-A similar case derived from an ecosystem example, involving declared types:
+And, in the current implementation, shadowing of module symbols (i.e., symbols exposed to other
+modules) cannot be recognized from lazy scopes.
 
 ```py
-class C: ...
-
-def outer(x: C | None):
-    x = x or C()
-
-    reveal_type(x)  # revealed: C
+class A: ...
+class A: ...
 
-    def inner() -> None:
-        # TODO: this should ideally be `C`
-        reveal_type(x)  # revealed: C | None
-    inner()
+def f(x: A):
+    # TODO: no error
+    # error: [invalid-assignment] "Object of type `A | A` is not assignable to `A`"
+    x = A()
 ```
 
 ### Assignments to nonlocal variables

crates/ty_python_semantic/src/semantic_index/builder.rs

@@ -115,6 +115,7 @@ pub(super) struct SemanticIndexBuilder<'db, 'ast> {
     ///
     /// [generator functions]: https://docs.python.org/3/glossary.html#term-generator
     generator_functions: FxHashSet<FileScopeId>,
+    /// Snapshots of enclosing-scope place states visible from nested scopes.
     enclosing_snapshots: FxHashMap<EnclosingSnapshotKey, ScopedEnclosingSnapshotId>,
     /// Errors collected by the `semantic_checker`.
     semantic_syntax_errors: RefCell<Vec<SemanticSyntaxError>>,
@@ -316,11 +317,12 @@ impl<'db, 'ast> SemanticIndexBuilder<'db, 'ast> {
                     nested_scope: popped_scope_id,
                     nested_laziness: ScopeLaziness::Eager,
                 };
-                let eager_snapshot = self.use_def_maps[enclosing_scope_id].snapshot_outer_state(
-                    enclosing_place_id,
-                    enclosing_scope_kind,
-                    enclosing_place,
-                );
+                let eager_snapshot = self.use_def_maps[enclosing_scope_id]
+                    .snapshot_enclosing_state(
+                        enclosing_place_id,
+                        enclosing_scope_kind,
+                        enclosing_place,
+                    );
                 self.enclosing_snapshots.insert(key, eager_snapshot);
             }
 
@@ -390,7 +392,7 @@ impl<'db, 'ast> SemanticIndexBuilder<'db, 'ast> {
                     nested_scope: popped_scope_id,
                     nested_laziness: ScopeLaziness::Lazy,
                 };
-                let lazy_snapshot = self.use_def_maps[enclosing_scope_id].snapshot_outer_state(
+                let lazy_snapshot = self.use_def_maps[enclosing_scope_id].snapshot_enclosing_state(
                     enclosed_symbol_id.into(),
                     enclosing_scope_kind,
                     enclosing_place.into(),
@@ -400,29 +402,74 @@ impl<'db, 'ast> SemanticIndexBuilder<'db, 'ast> {
         }
     }
 
-    /// Any lazy snapshots of places that have been reassigned are no longer valid, so delete them.
-    fn sweep_lazy_snapshots(&mut self, popped_scope_id: FileScopeId) {
-        // Retain only snapshots that are either eager
-        //     || (enclosing_scope != popped_scope && popped_scope is not a visible ancestor of enclosing_scope)
-        //     || enclosing_place is not a symbol or not reassigned
-        // <=> remove those that are lazy
-        //     && (enclosing_scope == popped_scope || popped_scope is a visible ancestor of enclosing_scope)
-        //     && enclosing_place is a symbol and reassigned
-        self.enclosing_snapshots.retain(|key, _| {
-            let popped_place_table = &self.place_tables[popped_scope_id];
-            key.nested_laziness.is_eager()
-                || (key.enclosing_scope != popped_scope_id
-                    && VisibleAncestorsIter::new(&self.scopes, key.enclosing_scope)
-                        .all(|(ancestor, _)| ancestor != popped_scope_id))
-                || !key.enclosing_place.as_symbol().is_some_and(|symbol_id| {
-                    let name = &self.place_tables[key.enclosing_scope]
-                        .symbol(symbol_id)
-                        .name();
-                    popped_place_table.symbol_id(name).is_some_and(|symbol_id| {
-                        popped_place_table.symbol(symbol_id).is_reassigned()
-                    })
-                })
-        });
+    /// Any lazy snapshots of the place that have been reassigned are obsolete, so update them.
+    /// ```py
+    /// def outer() -> None:
+    ///     x = None
+    ///
+    ///     def inner2() -> None:
+    ///         # `inner` can be referenced before its definition,
+    ///         # but `inner2` must still be called after the definition of `inner` for this call to be valid.
+    ///         inner()
+    ///
+    ///         # In this scope, `x` may refer to `x = None` or `x = 1`.
+    ///         reveal_type(x)  # revealed: None | Literal[1]
+    ///
+    ///     # Reassignment of `x` after the definition of `inner2`.
+    ///     # Update lazy snapshots of `x` for `inner2`.
+    ///     x = 1
+    ///
+    ///     def inner() -> None:
+    ///         # In this scope, `x = None` appears as being shadowed by `x = 1`.
+    ///         reveal_type(x)  # revealed: Literal[1]
+    ///
+    ///     # No reassignment of `x` after the definition of `inner`, so we can safely use a lazy snapshot for `inner` as is.
+    ///     inner()
+    ///     inner2()
+    /// ```
+    fn update_lazy_snapshots(&mut self, symbol: ScopedSymbolId) {
+        let current_scope = self.current_scope();
+        let current_place_table = &self.place_tables[current_scope];
+        let symbol = current_place_table.symbol(symbol);
+        // Optimization: if this is the first binding of the symbol we've seen, there can't be any
+        // lazy snapshots of it to update.
+        if !symbol.is_reassigned() {
+            return;
+        }
+        for (key, snapshot_id) in &self.enclosing_snapshots {
+            if let Some(enclosing_symbol) = key.enclosing_place.as_symbol() {
+                let name = self.place_tables[key.enclosing_scope]
+                    .symbol(enclosing_symbol)
+                    .name();
+                let is_reassignment_of_snapshotted_symbol = || {
+                    for (ancestor, _) in
+                        VisibleAncestorsIter::new(&self.scopes, key.enclosing_scope)
+                    {
+                        if ancestor == current_scope {
+                            return true;
+                        }
+                        let ancestor_table = &self.place_tables[ancestor];
+                        // If there is a symbol binding in an ancestor scope,
+                        // then a reassignment in the current scope is not relevant to the snapshot.
+                        if ancestor_table
+                            .symbol_id(name)
+                            .is_some_and(|id| ancestor_table.symbol(id).is_bound())
+                        {
+                            return false;
+                        }
+                    }
+                    false
+                };
+
+                if key.nested_laziness.is_lazy()
+                    && symbol.name() == name
+                    && is_reassignment_of_snapshotted_symbol()
+                {
+                    self.use_def_maps[key.enclosing_scope]
+                        .update_enclosing_snapshot(*snapshot_id, enclosing_symbol);
+                }
+            }
+        }
     }
 
     fn sweep_nonlocal_lazy_snapshots(&mut self) {
@@ -464,8 +511,6 @@ impl<'db, 'ast> SemanticIndexBuilder<'db, 'ast> {
             .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];
@@ -659,6 +704,12 @@ impl<'db, 'ast> SemanticIndexBuilder<'db, 'ast> {
             }
         }
 
+        if category.is_binding() {
+            if let Some(id) = place.as_symbol() {
+                self.update_lazy_snapshots(id);
+            }
+        }
+
         let mut try_node_stack_manager = std::mem::take(&mut self.try_node_context_stack_manager);
         try_node_stack_manager.record_definition(self);
         self.try_node_context_stack_manager = try_node_stack_manager;
@@ -1313,7 +1364,6 @@ impl<'ast> Visitor<'ast> for SemanticIndexBuilder<'_, 'ast> {
                 // used to collect all the overloaded definitions of a function. This needs to be
                 // done on the `Identifier` node as opposed to `ExprName` because that's what the
                 // AST uses.
-                self.mark_symbol_used(symbol);
                 let use_id = self.current_ast_ids().record_use(name);
                 self.current_use_def_map_mut().record_use(
                     symbol.into(),
@@ -1322,6 +1372,7 @@ impl<'ast> Visitor<'ast> for SemanticIndexBuilder<'_, 'ast> {
                 );
 
                 self.add_definition(symbol.into(), function_def);
+                self.mark_symbol_used(symbol);
             }
             ast::Stmt::ClassDef(class) => {
                 for decorator in &class.decorator_list {

crates/ty_python_semantic/src/semantic_index/scope.rs

@@ -195,6 +195,10 @@ impl ScopeLaziness {
     pub(crate) const fn is_eager(self) -> bool {
         matches!(self, ScopeLaziness::Eager)
     }
+
+    pub(crate) const fn is_lazy(self) -> bool {
+        matches!(self, ScopeLaziness::Lazy)
+    }
 }
 
 #[derive(Copy, Clone, Debug, PartialEq, Eq)]