[red-knot] Make Symbol::or_fall_back_to() lazy (#15943)

This commit is contained in:
Alex Waygood 2025-02-05 14:51:02 +00:00 committed by GitHub
parent c69b19fe1d
commit 2ebb5e8d4b
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
4 changed files with 99 additions and 100 deletions

View file

@ -13,7 +13,7 @@ if returns_bool():
chr: int = 1
def f():
reveal_type(chr) # revealed: Literal[chr] | int
reveal_type(chr) # revealed: int | Literal[chr]
```
## Conditionally global or builtin, with annotation
@ -28,5 +28,5 @@ if returns_bool():
chr: int = 1
def f():
reveal_type(chr) # revealed: Literal[chr] | int
reveal_type(chr) # revealed: int | Literal[chr]
```

View file

@ -9,15 +9,6 @@ pub(crate) enum Boundness {
PossiblyUnbound,
}
impl Boundness {
pub(crate) fn or(self, other: Boundness) -> Boundness {
match (self, other) {
(Boundness::Bound, _) | (_, Boundness::Bound) => Boundness::Bound,
(Boundness::PossiblyUnbound, Boundness::PossiblyUnbound) => Boundness::PossiblyUnbound,
}
}
}
/// The result of a symbol lookup, which can either be a (possibly unbound) type
/// or a completely unbound symbol.
///
@ -46,13 +37,6 @@ impl<'db> Symbol<'db> {
matches!(self, Symbol::Unbound)
}
pub(crate) fn possibly_unbound(&self) -> bool {
match self {
Symbol::Type(_, Boundness::PossiblyUnbound) | Symbol::Unbound => true,
Symbol::Type(_, Boundness::Bound) => false,
}
}
/// Returns the type of the symbol, ignoring possible unboundness.
///
/// If the symbol is *definitely* unbound, this function will return `None`. Otherwise,
@ -71,18 +55,32 @@ impl<'db> Symbol<'db> {
.expect("Expected a (possibly unbound) type, not an unbound symbol")
}
/// Fallback (partially or fully) to another symbol if `self` is partially or fully unbound.
///
/// 1. If `self` is definitely bound, return `self` without evaluating `fallback_fn()`.
/// 2. Else, evaluate `fallback_fn()`:
/// a. If `self` is definitely unbound, return the result of `fallback_fn()`.
/// b. Else, if `fallback` is definitely unbound, return `self`.
/// c. Else, if `self` is possibly unbound and `fallback` is definitely bound,
/// return `Symbol(<union of self-type and fallback-type>, Boundness::Bound)`
/// d. Else, if `self` is possibly unbound and `fallback` is possibly unbound,
/// return `Symbol(<union of self-type and fallback-type>, Boundness::PossiblyUnbound)`
#[must_use]
pub(crate) fn or_fall_back_to(self, db: &'db dyn Db, fallback: &Symbol<'db>) -> Symbol<'db> {
match fallback {
Symbol::Type(fallback_ty, fallback_boundness) => match self {
Symbol::Type(_, Boundness::Bound) => self,
Symbol::Type(ty, boundness @ Boundness::PossiblyUnbound) => Symbol::Type(
UnionType::from_elements(db, [*fallback_ty, ty]),
fallback_boundness.or(boundness),
pub(crate) fn or_fall_back_to(
self,
db: &'db dyn Db,
fallback_fn: impl FnOnce() -> Self,
) -> Self {
match self {
Symbol::Type(_, Boundness::Bound) => self,
Symbol::Unbound => fallback_fn(),
Symbol::Type(self_ty, Boundness::PossiblyUnbound) => match fallback_fn() {
Symbol::Unbound => self,
Symbol::Type(fallback_ty, fallback_boundness) => Symbol::Type(
UnionType::from_elements(db, [self_ty, fallback_ty]),
fallback_boundness,
),
Symbol::Unbound => fallback.clone(),
},
Symbol::Unbound => self,
}
}
@ -110,44 +108,44 @@ mod tests {
// Start from an unbound symbol
assert_eq!(
Symbol::Unbound.or_fall_back_to(&db, &Symbol::Unbound),
Symbol::Unbound.or_fall_back_to(&db, || Symbol::Unbound),
Symbol::Unbound
);
assert_eq!(
Symbol::Unbound.or_fall_back_to(&db, &Symbol::Type(ty1, PossiblyUnbound)),
Symbol::Unbound.or_fall_back_to(&db, || Symbol::Type(ty1, PossiblyUnbound)),
Symbol::Type(ty1, PossiblyUnbound)
);
assert_eq!(
Symbol::Unbound.or_fall_back_to(&db, &Symbol::Type(ty1, Bound)),
Symbol::Unbound.or_fall_back_to(&db, || Symbol::Type(ty1, Bound)),
Symbol::Type(ty1, Bound)
);
// Start from a possibly unbound symbol
assert_eq!(
Symbol::Type(ty1, PossiblyUnbound).or_fall_back_to(&db, &Symbol::Unbound),
Symbol::Type(ty1, PossiblyUnbound).or_fall_back_to(&db, || Symbol::Unbound),
Symbol::Type(ty1, PossiblyUnbound)
);
assert_eq!(
Symbol::Type(ty1, PossiblyUnbound)
.or_fall_back_to(&db, &Symbol::Type(ty2, PossiblyUnbound)),
Symbol::Type(UnionType::from_elements(&db, [ty2, ty1]), PossiblyUnbound)
.or_fall_back_to(&db, || Symbol::Type(ty2, PossiblyUnbound)),
Symbol::Type(UnionType::from_elements(&db, [ty1, ty2]), PossiblyUnbound)
);
assert_eq!(
Symbol::Type(ty1, PossiblyUnbound).or_fall_back_to(&db, &Symbol::Type(ty2, Bound)),
Symbol::Type(UnionType::from_elements(&db, [ty2, ty1]), Bound)
Symbol::Type(ty1, PossiblyUnbound).or_fall_back_to(&db, || Symbol::Type(ty2, Bound)),
Symbol::Type(UnionType::from_elements(&db, [ty1, ty2]), Bound)
);
// Start from a definitely bound symbol
assert_eq!(
Symbol::Type(ty1, Bound).or_fall_back_to(&db, &Symbol::Unbound),
Symbol::Type(ty1, Bound).or_fall_back_to(&db, || Symbol::Unbound),
Symbol::Type(ty1, Bound)
);
assert_eq!(
Symbol::Type(ty1, Bound).or_fall_back_to(&db, &Symbol::Type(ty2, PossiblyUnbound)),
Symbol::Type(ty1, Bound).or_fall_back_to(&db, || Symbol::Type(ty2, PossiblyUnbound)),
Symbol::Type(ty1, Bound)
);
assert_eq!(
Symbol::Type(ty1, Bound).or_fall_back_to(&db, &Symbol::Type(ty2, Bound)),
Symbol::Type(ty1, Bound).or_fall_back_to(&db, || Symbol::Type(ty2, Bound)),
Symbol::Type(ty1, Bound)
);
}

View file

@ -256,26 +256,21 @@ fn module_type_symbols<'db>(db: &'db dyn Db) -> smallvec::SmallVec<[ast::name::N
/// Looks up a module-global symbol by name in a file.
pub(crate) fn global_symbol<'db>(db: &'db dyn Db, file: File, name: &str) -> Symbol<'db> {
let explicit_symbol = symbol(db, global_scope(db, file), name);
if !explicit_symbol.possibly_unbound() {
return explicit_symbol;
}
// Not defined explicitly in the global scope?
// All modules are instances of `types.ModuleType`;
// look it up there (with a few very special exceptions)
if module_type_symbols(db)
.iter()
.any(|module_type_member| &**module_type_member == name)
{
// TODO: this should use `.to_instance(db)`. but we don't understand attribute access
// on instance types yet.
let module_type_member = KnownClass::ModuleType.to_class_literal(db).member(db, name);
return explicit_symbol.or_fall_back_to(db, &module_type_member);
}
explicit_symbol
symbol(db, global_scope(db, file), name).or_fall_back_to(db, || {
if module_type_symbols(db)
.iter()
.any(|module_type_member| &**module_type_member == name)
{
// TODO: this should use `.to_instance(db)`. but we don't understand attribute access
// on instance types yet.
KnownClass::ModuleType.to_class_literal(db).member(db, name)
} else {
Symbol::Unbound
}
})
}
/// Infer the type of a binding.
@ -3796,10 +3791,8 @@ impl<'db> ModuleLiteralType<'db> {
}
}
let global_lookup = symbol(db, global_scope(db, self.module(db).file()), name);
// If it's unbound, check if it's present as an instance on `types.ModuleType`
// or `builtins.object`.
// If it's not found in the global scope, check if it's present as an instance
// on `types.ModuleType` or `builtins.object`.
//
// We do a more limited version of this in `global_symbol_ty`,
// but there are two crucial differences here:
@ -3813,14 +3806,14 @@ impl<'db> ModuleLiteralType<'db> {
// ignore `__getattr__`. Typeshed has a fake `__getattr__` on `types.ModuleType`
// to help out with dynamic imports; we shouldn't use it for `ModuleLiteral` types
// where we know exactly which module we're dealing with.
if name != "__getattr__" && global_lookup.possibly_unbound() {
// TODO: this should use `.to_instance()`, but we don't understand instance attribute yet
let module_type_instance_member =
KnownClass::ModuleType.to_class_literal(db).member(db, name);
global_lookup.or_fall_back_to(db, &module_type_instance_member)
} else {
global_lookup
}
symbol(db, global_scope(db, self.module(db).file()), name).or_fall_back_to(db, || {
if name == "__getattr__" {
Symbol::Unbound
} else {
// TODO: this should use `.to_instance()`, but we don't understand instance attribute yet
KnownClass::ModuleType.to_class_literal(db).member(db, name)
}
})
}
}

View file

@ -3290,8 +3290,9 @@ impl<'db> TypeInferenceBuilder<'db> {
/// Look up a name reference that isn't bound in the local scope.
fn lookup_name(&mut self, name_node: &ast::ExprName) -> Symbol<'db> {
let db = self.db();
let ast::ExprName { id: name, .. } = name_node;
let file_scope_id = self.scope().file_scope_id(self.db());
let file_scope_id = self.scope().file_scope_id(db);
let is_bound =
if let Some(symbol) = self.index.symbol_table(file_scope_id).symbol_by_name(name) {
symbol.is_bound()
@ -3306,16 +3307,15 @@ impl<'db> TypeInferenceBuilder<'db> {
// In function-like scopes, any local variable (symbol that is bound in this scope) can
// only have a definition in this scope, or error; it never references another scope.
// (At runtime, it would use the `LOAD_FAST` opcode.)
if !is_bound || !self.scope().is_function_like(self.db()) {
if !is_bound || !self.scope().is_function_like(db) {
// Walk up parent scopes looking for a possible enclosing scope that may have a
// definition of this name visible to us (would be `LOAD_DEREF` at runtime.)
for (enclosing_scope_file_id, _) in self.index.ancestor_scopes(file_scope_id) {
// Class scopes are not visible to nested scopes, and we need to handle global
// scope differently (because an unbound name there falls back to builtins), so
// check only function-like scopes.
let enclosing_scope_id =
enclosing_scope_file_id.to_scope_id(self.db(), self.file());
if !enclosing_scope_id.is_function_like(self.db()) {
let enclosing_scope_id = enclosing_scope_file_id.to_scope_id(db, self.file());
if !enclosing_scope_id.is_function_like(db) {
continue;
}
let enclosing_symbol_table = self.index.symbol_table(enclosing_scope_file_id);
@ -3328,37 +3328,45 @@ impl<'db> TypeInferenceBuilder<'db> {
// runtime, it is the scope that creates the cell for our closure.) If the name
// isn't bound in that scope, we should get an unbound name, not continue
// falling back to other scopes / globals / builtins.
return symbol(self.db(), enclosing_scope_id, name);
return symbol(db, enclosing_scope_id, name);
}
}
// No nonlocal binding, check module globals. Avoid infinite recursion if `self.scope`
// already is module globals.
let global_symbol = if file_scope_id.is_global() {
Symbol::Unbound
} else {
global_symbol(self.db(), self.file(), name)
};
// Fallback to builtins (without infinite recursion if we're already in builtins.)
if global_symbol.possibly_unbound()
&& Some(self.scope()) != builtins_module_scope(self.db())
{
let mut builtins_symbol = builtins_symbol(self.db(), name);
if builtins_symbol.is_unbound() && name == "reveal_type" {
self.context.report_lint(
&UNDEFINED_REVEAL,
name_node.into(),
format_args!(
"`reveal_type` used without importing it; this is allowed for debugging convenience but will fail at runtime"),
);
builtins_symbol = typing_extensions_symbol(self.db(), name);
}
global_symbol.or_fall_back_to(self.db(), &builtins_symbol)
} else {
global_symbol
}
Symbol::Unbound
// No nonlocal binding? Check the module's globals.
// Avoid infinite recursion if `self.scope` already is the module's global scope.
.or_fall_back_to(db, || {
if file_scope_id.is_global() {
Symbol::Unbound
} else {
global_symbol(db, self.file(), name)
}
})
// Not found in globals? Fallback to builtins
// (without infinite recursion if we're already in builtins.)
.or_fall_back_to(db, || {
if Some(self.scope()) == builtins_module_scope(db) {
Symbol::Unbound
} else {
builtins_symbol(db, name)
}
})
// Still not found? It might be `reveal_type`...
.or_fall_back_to(db, || {
if name == "reveal_type" {
self.context.report_lint(
&UNDEFINED_REVEAL,
name_node.into(),
format_args!(
"`reveal_type` used without importing it; \
this is allowed for debugging convenience but will fail at runtime"
),
);
typing_extensions_symbol(db, name)
} else {
Symbol::Unbound
}
})
} else {
Symbol::Unbound
}