[ty] more precise lazy scope place lookup (#19932)

## Summary

This is a follow-up to https://github.com/astral-sh/ruff/pull/19321.

Now lazy snapshots are updated to take into account new bindings on
every symbol reassignment.

```python
def outer(x: A | None):
    if x is None:
        x = A()

    reveal_type(x)  # revealed: A

    def inner() -> None:
        # lazy snapshot: {x: A}
        reveal_type(x)  # revealed: A
    inner()

def outer() -> None:
    x = None

    x = 1

    def inner() -> None:
        # lazy snapshot: {x: Literal[1]} -> {x: Literal[1, 2]}
        reveal_type(x)  # revealed: Literal[1, 2]
    inner()

    x = 2
```

Closes astral-sh/ty#559.

## Test Plan

Some TODOs in `public_types.md` now work properly.

---------

Co-authored-by: Carl Meyer <carl@astral.sh>
This commit is contained in:
Shunsuke Shibayama 2025-09-09 06:08:35 +09:00 committed by GitHub
parent aa5d665d52
commit 08a561fc05
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
7 changed files with 208 additions and 84 deletions

View file

@ -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]):

View file

@ -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`:
And, in the current implementation, shadowing of module symbols (i.e., symbols exposed to other
modules) cannot be recognized from lazy scopes.
```py
def outer() -> None:
x = None
class A: ...
class A: ...
# [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:
```py
class C: ...
def outer(x: C | None):
x = x or C()
reveal_type(x) # revealed: C
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

View file

@ -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,7 +317,8 @@ 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(
let eager_snapshot = self.use_def_maps[enclosing_scope_id]
.snapshot_enclosing_state(
enclosing_place_id,
enclosing_scope_kind,
enclosing_place,
@ -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)
/// 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();
popped_place_table.symbol_id(name).is_some_and(|symbol_id| {
popped_place_table.symbol(symbol_id).is_reassigned()
})
})
});
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 {

View file

@ -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)]

View file

@ -243,10 +243,6 @@
use ruff_index::{IndexVec, newtype_index};
use rustc_hash::FxHashMap;
use self::place_state::{
Bindings, Declarations, EnclosingSnapshot, LiveBindingsIterator, LiveDeclaration,
LiveDeclarationsIterator, PlaceState, ScopedDefinitionId,
};
use crate::node_key::NodeKey;
use crate::place::BoundnessAnalysis;
use crate::semantic_index::ast_ids::ScopedUseId;
@ -254,6 +250,7 @@ use crate::semantic_index::definition::{Definition, DefinitionState};
use crate::semantic_index::member::ScopedMemberId;
use crate::semantic_index::narrowing_constraints::{
ConstraintKey, NarrowingConstraints, NarrowingConstraintsBuilder, NarrowingConstraintsIterator,
ScopedNarrowingConstraint,
};
use crate::semantic_index::place::{PlaceExprRef, ScopedPlaceId};
use crate::semantic_index::predicate::{
@ -264,7 +261,10 @@ use crate::semantic_index::reachability_constraints::{
};
use crate::semantic_index::scope::{FileScopeId, ScopeKind, ScopeLaziness};
use crate::semantic_index::symbol::ScopedSymbolId;
use crate::semantic_index::use_def::place_state::PreviousDefinitions;
use crate::semantic_index::use_def::place_state::{
Bindings, Declarations, EnclosingSnapshot, LiveBindingsIterator, LiveDeclaration,
LiveDeclarationsIterator, PlaceState, PreviousDefinitions, ScopedDefinitionId,
};
use crate::semantic_index::{EnclosingSnapshotResult, SemanticIndex};
use crate::types::{IntersectionBuilder, Truthiness, Type, infer_narrowing_constraint};
@ -640,8 +640,8 @@ impl<'db> UseDefMap<'db> {
}
}
/// Uniquely identifies a snapshot of an enclosing scope place state that can be used to resolve a reference in a
/// nested scope.
/// Uniquely identifies a snapshot of an enclosing scope place state that can be used to resolve a
/// reference in a 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
@ -660,16 +660,14 @@ pub(crate) struct EnclosingSnapshotKey {
pub(crate) enclosing_place: ScopedPlaceId,
/// The nested scope containing the reference
pub(crate) nested_scope: FileScopeId,
/// Laziness of the nested scope
/// Laziness of the nested scope (technically redundant, but convenient to have here)
pub(crate) nested_laziness: ScopeLaziness,
}
/// A snapshot of enclosing scope place states that can be used to resolve a reference in a nested scope.
/// 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.
/// Snapshots of enclosing scope place states for resolving a reference in a nested scope.
/// If the nested scope is eager, the snapshot is simply recorded and used as is.
/// If it is lazy, every time the outer symbol is reassigned, the snapshot is updated to add the
/// new binding.
type EnclosingSnapshots = IndexVec<ScopedEnclosingSnapshotId, EnclosingSnapshot>;
#[derive(Debug)]
@ -1185,7 +1183,7 @@ impl<'db> UseDefMapBuilder<'db> {
self.node_reachability.insert(node_key, self.reachability);
}
pub(super) fn snapshot_outer_state(
pub(super) fn snapshot_enclosing_state(
&mut self,
enclosing_place: ScopedPlaceId,
scope: ScopeKind,
@ -1208,6 +1206,27 @@ impl<'db> UseDefMapBuilder<'db> {
}
}
pub(super) fn update_enclosing_snapshot(
&mut self,
snapshot_id: ScopedEnclosingSnapshotId,
enclosing_symbol: ScopedSymbolId,
) {
match self.enclosing_snapshots.get_mut(snapshot_id) {
Some(EnclosingSnapshot::Bindings(bindings)) => {
let new_symbol_state = &self.symbol_states[enclosing_symbol];
bindings.merge(
new_symbol_state.bindings().clone(),
&mut self.narrowing_constraints,
&mut self.reachability_constraints,
);
}
Some(EnclosingSnapshot::Constraint(constraint)) => {
*constraint = ScopedNarrowingConstraint::empty();
}
None => {}
}
}
/// Take a snapshot of the current visible-places state.
pub(super) fn snapshot(&self) -> FlowSnapshot {
FlowSnapshot {

View file

@ -308,7 +308,7 @@ impl Bindings {
self.live_bindings.iter()
}
fn merge(
pub(super) fn merge(
&mut self,
b: Self,
narrowing_constraints: &mut NarrowingConstraintsBuilder,

View file

@ -6936,12 +6936,6 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> {
));
}
EnclosingSnapshotResult::FoundBindings(bindings) => {
if place_expr.is_symbol()
&& !enclosing_scope_id.is_function_like(db)
&& !is_immediately_enclosing_scope
{
continue;
}
let place = place_from_bindings(db, bindings).map_type(|ty| {
self.narrow_place_with_applicable_constraints(
place_expr,