[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 chr: int = 1
def f(): def f():
reveal_type(chr) # revealed: Literal[chr] | int reveal_type(chr) # revealed: int | Literal[chr]
``` ```
## Conditionally global or builtin, with annotation ## Conditionally global or builtin, with annotation
@ -28,5 +28,5 @@ if returns_bool():
chr: int = 1 chr: int = 1
def f(): 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, 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 /// The result of a symbol lookup, which can either be a (possibly unbound) type
/// or a completely unbound symbol. /// or a completely unbound symbol.
/// ///
@ -46,13 +37,6 @@ impl<'db> Symbol<'db> {
matches!(self, Symbol::Unbound) 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. /// Returns the type of the symbol, ignoring possible unboundness.
/// ///
/// If the symbol is *definitely* unbound, this function will return `None`. Otherwise, /// 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") .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] #[must_use]
pub(crate) fn or_fall_back_to(self, db: &'db dyn Db, fallback: &Symbol<'db>) -> Symbol<'db> { pub(crate) fn or_fall_back_to(
match fallback { self,
Symbol::Type(fallback_ty, fallback_boundness) => match self { db: &'db dyn Db,
Symbol::Type(_, Boundness::Bound) => self, fallback_fn: impl FnOnce() -> Self,
Symbol::Type(ty, boundness @ Boundness::PossiblyUnbound) => Symbol::Type( ) -> Self {
UnionType::from_elements(db, [*fallback_ty, ty]), match self {
fallback_boundness.or(boundness), 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 // Start from an unbound symbol
assert_eq!( assert_eq!(
Symbol::Unbound.or_fall_back_to(&db, &Symbol::Unbound), Symbol::Unbound.or_fall_back_to(&db, || Symbol::Unbound),
Symbol::Unbound Symbol::Unbound
); );
assert_eq!( 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) Symbol::Type(ty1, PossiblyUnbound)
); );
assert_eq!( 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) Symbol::Type(ty1, Bound)
); );
// Start from a possibly unbound symbol // Start from a possibly unbound symbol
assert_eq!( 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) Symbol::Type(ty1, PossiblyUnbound)
); );
assert_eq!( assert_eq!(
Symbol::Type(ty1, PossiblyUnbound) Symbol::Type(ty1, PossiblyUnbound)
.or_fall_back_to(&db, &Symbol::Type(ty2, PossiblyUnbound)), .or_fall_back_to(&db, || Symbol::Type(ty2, PossiblyUnbound)),
Symbol::Type(UnionType::from_elements(&db, [ty2, ty1]), PossiblyUnbound) Symbol::Type(UnionType::from_elements(&db, [ty1, ty2]), PossiblyUnbound)
); );
assert_eq!( assert_eq!(
Symbol::Type(ty1, PossiblyUnbound).or_fall_back_to(&db, &Symbol::Type(ty2, Bound)), Symbol::Type(ty1, PossiblyUnbound).or_fall_back_to(&db, || Symbol::Type(ty2, Bound)),
Symbol::Type(UnionType::from_elements(&db, [ty2, ty1]), Bound) Symbol::Type(UnionType::from_elements(&db, [ty1, ty2]), Bound)
); );
// Start from a definitely bound symbol // Start from a definitely bound symbol
assert_eq!( 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) Symbol::Type(ty1, Bound)
); );
assert_eq!( 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) Symbol::Type(ty1, Bound)
); );
assert_eq!( 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) 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. /// 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> { 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? // Not defined explicitly in the global scope?
// All modules are instances of `types.ModuleType`; // All modules are instances of `types.ModuleType`;
// look it up there (with a few very special exceptions) // look it up there (with a few very special exceptions)
if module_type_symbols(db) symbol(db, global_scope(db, file), name).or_fall_back_to(db, || {
.iter() if module_type_symbols(db)
.any(|module_type_member| &**module_type_member == name) .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. // TODO: this should use `.to_instance(db)`. but we don't understand attribute access
let module_type_member = KnownClass::ModuleType.to_class_literal(db).member(db, name); // on instance types yet.
return explicit_symbol.or_fall_back_to(db, &module_type_member); KnownClass::ModuleType.to_class_literal(db).member(db, name)
} } else {
Symbol::Unbound
explicit_symbol }
})
} }
/// Infer the type of a binding. /// 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 not found in the global scope, check if it's present as an instance
// on `types.ModuleType` or `builtins.object`.
// If it's unbound, 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`, // We do a more limited version of this in `global_symbol_ty`,
// but there are two crucial differences here: // 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` // ignore `__getattr__`. Typeshed has a fake `__getattr__` on `types.ModuleType`
// to help out with dynamic imports; we shouldn't use it for `ModuleLiteral` types // to help out with dynamic imports; we shouldn't use it for `ModuleLiteral` types
// where we know exactly which module we're dealing with. // where we know exactly which module we're dealing with.
if name != "__getattr__" && global_lookup.possibly_unbound() { symbol(db, global_scope(db, self.module(db).file()), name).or_fall_back_to(db, || {
// TODO: this should use `.to_instance()`, but we don't understand instance attribute yet if name == "__getattr__" {
let module_type_instance_member = Symbol::Unbound
KnownClass::ModuleType.to_class_literal(db).member(db, name); } else {
global_lookup.or_fall_back_to(db, &module_type_instance_member) // TODO: this should use `.to_instance()`, but we don't understand instance attribute yet
} else { KnownClass::ModuleType.to_class_literal(db).member(db, name)
global_lookup }
} })
} }
} }

View file

@ -3290,8 +3290,9 @@ impl<'db> TypeInferenceBuilder<'db> {
/// Look up a name reference that isn't bound in the local scope. /// Look up a name reference that isn't bound in the local scope.
fn lookup_name(&mut self, name_node: &ast::ExprName) -> Symbol<'db> { fn lookup_name(&mut self, name_node: &ast::ExprName) -> Symbol<'db> {
let db = self.db();
let ast::ExprName { id: name, .. } = name_node; 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 = let is_bound =
if let Some(symbol) = self.index.symbol_table(file_scope_id).symbol_by_name(name) { if let Some(symbol) = self.index.symbol_table(file_scope_id).symbol_by_name(name) {
symbol.is_bound() 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 // 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. // only have a definition in this scope, or error; it never references another scope.
// (At runtime, it would use the `LOAD_FAST` opcode.) // (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 // 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.) // 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) { 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 // 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 // scope differently (because an unbound name there falls back to builtins), so
// check only function-like scopes. // check only function-like scopes.
let enclosing_scope_id = let enclosing_scope_id = enclosing_scope_file_id.to_scope_id(db, self.file());
enclosing_scope_file_id.to_scope_id(self.db(), self.file()); if !enclosing_scope_id.is_function_like(db) {
if !enclosing_scope_id.is_function_like(self.db()) {
continue; continue;
} }
let enclosing_symbol_table = self.index.symbol_table(enclosing_scope_file_id); 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 // 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 // isn't bound in that scope, we should get an unbound name, not continue
// falling back to other scopes / globals / builtins. // 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` Symbol::Unbound
// already is module globals. // No nonlocal binding? Check the module's globals.
let global_symbol = if file_scope_id.is_global() { // Avoid infinite recursion if `self.scope` already is the module's global scope.
Symbol::Unbound .or_fall_back_to(db, || {
} else { if file_scope_id.is_global() {
global_symbol(self.db(), self.file(), name) Symbol::Unbound
}; } else {
global_symbol(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()) // Not found in globals? Fallback to builtins
{ // (without infinite recursion if we're already in builtins.)
let mut builtins_symbol = builtins_symbol(self.db(), name); .or_fall_back_to(db, || {
if builtins_symbol.is_unbound() && name == "reveal_type" { if Some(self.scope()) == builtins_module_scope(db) {
self.context.report_lint( Symbol::Unbound
&UNDEFINED_REVEAL, } else {
name_node.into(), builtins_symbol(db, name)
format_args!( }
"`reveal_type` used without importing it; this is allowed for debugging convenience but will fail at runtime"), })
); // Still not found? It might be `reveal_type`...
builtins_symbol = typing_extensions_symbol(self.db(), name); .or_fall_back_to(db, || {
} if name == "reveal_type" {
self.context.report_lint(
global_symbol.or_fall_back_to(self.db(), &builtins_symbol) &UNDEFINED_REVEAL,
} else { name_node.into(),
global_symbol 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 { } else {
Symbol::Unbound Symbol::Unbound
} }