[red-knot] Symbol API improvements, part 2 (#14276)

## Summary

Apart from one small functional change, this is mostly a refactoring of
the `Symbol` API:

- Rename `as_type` to the more explicit `ignore_possibly_unbound`, no
functional change
- Remove `unwrap_or_unknown` in favor of the more explicit
`.ignore_possibly_unbound().unwrap_or(Type::Unknown)`, no functional
change
- Consistently call it "possibly unbound" (not "may be unbound")
- Rename `replace_unbound_with` to `or_fall_back_to` and properly handle
boundness of the fall back. This is the only functional change (did not
have any impact on existing tests).

relates to: #14022

## Test Plan

New unit tests for `Symbol::or_fall_back_to`

---------

Co-authored-by: Alex Waygood <Alex.Waygood@Gmail.com>

crates/red_knot_python_semantic/src/semantic_index/use_def.rs

@@ -277,7 +277,7 @@ impl<'db> UseDefMap<'db> {
 
     pub(crate) fn use_boundness(&self, use_id: ScopedUseId) -> Boundness {
         if self.bindings_by_use[use_id].may_be_unbound() {
-            Boundness::MayBeUnbound
+            Boundness::PossiblyUnbound
         } else {
             Boundness::Bound
         }
@@ -292,7 +292,7 @@ impl<'db> UseDefMap<'db> {
 
     pub(crate) fn public_boundness(&self, symbol: ScopedSymbolId) -> Boundness {
         if self.public_symbols[symbol].may_be_unbound() {
-            Boundness::MayBeUnbound
+            Boundness::PossiblyUnbound
         } else {
             Boundness::Bound
         }

crates/red_knot_python_semantic/src/symbol.rs

@@ -6,7 +6,16 @@ use crate::{
 #[derive(Debug, Clone, Copy, PartialEq, Eq)]
 pub(crate) enum Boundness {
     Bound,
-    MayBeUnbound,
+    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
@@ -17,14 +26,14 @@ pub(crate) enum Boundness {
 /// bound = 1
 ///
 /// if flag:
-///     maybe_unbound = 2
+///     possibly_unbound = 2
 /// ```
 ///
 /// If we look up symbols in this scope, we would get the following results:
 /// ```rs
-/// bound:          Symbol::Type(Type::IntLiteral(1), Boundness::Bound),
-/// maybe_unbound:  Symbol::Type(Type::IntLiteral(2), Boundness::MayBeUnbound),
-/// non_existent:   Symbol::Unbound,
+/// bound:             Symbol::Type(Type::IntLiteral(1), Boundness::Bound),
+/// possibly_unbound:  Symbol::Type(Type::IntLiteral(2), Boundness::PossiblyUnbound),
+/// non_existent:      Symbol::Unbound,
 /// ```
 #[derive(Debug, Clone, PartialEq)]
 pub(crate) enum Symbol<'db> {
@@ -37,21 +46,18 @@ impl<'db> Symbol<'db> {
         matches!(self, Symbol::Unbound)
     }
 
-    pub(crate) fn may_be_unbound(&self) -> bool {
+    pub(crate) fn possibly_unbound(&self) -> bool {
         match self {
-            Symbol::Type(_, Boundness::MayBeUnbound) | Symbol::Unbound => true,
+            Symbol::Type(_, Boundness::PossiblyUnbound) | Symbol::Unbound => true,
             Symbol::Type(_, Boundness::Bound) => false,
         }
     }
 
-    pub(crate) fn unwrap_or_unknown(&self) -> Type<'db> {
-        match self {
-            Symbol::Type(ty, _) => *ty,
-            Symbol::Unbound => Type::Unknown,
-        }
-    }
-
-    pub(crate) fn as_type(&self) -> Option<Type<'db>> {
+    /// Returns the type of the symbol, ignoring possible unboundness.
+    ///
+    /// If the symbol is *definitely* unbound, this function will return `None`. Otherwise,
+    /// if there is at least one control-flow path where the symbol is bound, return the type.
+    pub(crate) fn ignore_possibly_unbound(&self) -> Option<Type<'db>> {
         match self {
             Symbol::Type(ty, _) => Some(*ty),
             Symbol::Unbound => None,
@@ -61,28 +67,80 @@ impl<'db> Symbol<'db> {
     #[cfg(test)]
     #[track_caller]
     pub(crate) fn expect_type(self) -> Type<'db> {
-        self.as_type()
+        self.ignore_possibly_unbound()
             .expect("Expected a (possibly unbound) type, not an unbound symbol")
     }
 
     #[must_use]
-    pub(crate) fn replace_unbound_with(
-        self,
-        db: &'db dyn Db,
-        replacement: &Symbol<'db>,
-    ) -> Symbol<'db> {
-        match replacement {
-            Symbol::Type(replacement, _) => Symbol::Type(
-                match self {
-                    Symbol::Type(ty, Boundness::Bound) => ty,
-                    Symbol::Type(ty, Boundness::MayBeUnbound) => {
-                        UnionType::from_elements(db, [*replacement, ty])
-                    }
-                    Symbol::Unbound => *replacement,
-                },
-                Boundness::Bound,
-            ),
+    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),
+                ),
+                Symbol::Unbound => fallback.clone(),
+            },
             Symbol::Unbound => self,
         }
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use crate::types::tests::setup_db;
+
+    #[test]
+    fn test_symbol_or_fall_back_to() {
+        use Boundness::{Bound, PossiblyUnbound};
+
+        let db = setup_db();
+        let ty1 = Type::IntLiteral(1);
+        let ty2 = Type::IntLiteral(2);
+
+        // Start from an unbound symbol
+        assert_eq!(
+            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::Type(ty1, PossiblyUnbound)
+        );
+        assert_eq!(
+            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)
+        );
+        assert_eq!(
+            Symbol::Type(ty1, PossiblyUnbound)
+                .or_fall_back_to(&db, &Symbol::Type(ty2, PossiblyUnbound)),
+            Symbol::Type(UnionType::from_elements(&db, [ty2, ty1]), 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)
+        );
+
+        // Start from a definitely bound symbol
+        assert_eq!(
+            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)
+        );
+        assert_eq!(
+            Symbol::Type(ty1, Bound).or_fall_back_to(&db, &Symbol::Type(ty2, Bound)),
+            Symbol::Type(ty1, Bound)
+        );
+    }
+}

crates/red_knot_python_semantic/src/types.rs

@@ -157,7 +157,7 @@ fn module_type_symbols<'db>(db: &'db dyn Db) -> smallvec::SmallVec<[ast::name::N
 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.may_be_unbound() {
+    if !explicit_symbol.possibly_unbound() {
         return explicit_symbol;
     }
 
@@ -171,7 +171,7 @@ pub(crate) fn global_symbol<'db>(db: &'db dyn Db, file: File, name: &str) -> Sym
         // 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.replace_unbound_with(db, &module_type_member);
+        return explicit_symbol.or_fall_back_to(db, &module_type_member);
     }
 
     explicit_symbol
@@ -1071,11 +1071,11 @@ impl<'db> Type<'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.may_be_unbound() {
+                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.replace_unbound_with(db, &module_type_instance_member)
+                    global_lookup.or_fall_back_to(db, &module_type_instance_member)
                 } else {
                     global_lookup
                 }
@@ -1100,17 +1100,17 @@ impl<'db> Type<'db> {
                 let mut builder = UnionBuilder::new(db);
 
                 let mut all_unbound = true;
-                let mut may_be_unbound = false;
+                let mut possibly_unbound = false;
                 for ty in union.elements(db) {
                     let ty_member = ty.member(db, name);
                     match ty_member {
                         Symbol::Unbound => {
-                            may_be_unbound = true;
+                            possibly_unbound = true;
                         }
                         Symbol::Type(ty_member, member_boundness) => {
                             // TODO: raise a diagnostic if member_boundness indicates potential unboundness
-                            if member_boundness == Boundness::MayBeUnbound {
-                                may_be_unbound = true;
+                            if member_boundness == Boundness::PossiblyUnbound {
+                                possibly_unbound = true;
                             }
 
                             all_unbound = false;
@@ -1124,8 +1124,8 @@ impl<'db> Type<'db> {
                 } else {
                     Symbol::Type(
                         builder.build(),
-                        if may_be_unbound {
-                            Boundness::MayBeUnbound
+                        if possibly_unbound {
+                            Boundness::PossiblyUnbound
                         } else {
                             Boundness::Bound
                         },
@@ -1319,7 +1319,7 @@ impl<'db> Type<'db> {
             Symbol::Type(callable_ty, Boundness::Bound) => {
                 CallDunderResult::CallOutcome(callable_ty.call(db, arg_types))
             }
-            Symbol::Type(callable_ty, Boundness::MayBeUnbound) => {
+            Symbol::Type(callable_ty, Boundness::PossiblyUnbound) => {
                 CallDunderResult::PossiblyUnbound(callable_ty.call(db, arg_types))
             }
             Symbol::Unbound => CallDunderResult::MethodNotAvailable,
@@ -1625,7 +1625,9 @@ impl<'db> KnownClass {
     }
 
     pub fn to_class_literal(self, db: &'db dyn Db) -> Type<'db> {
-        core_module_symbol(db, self.canonical_module(), self.as_str()).unwrap_or_unknown()
+        core_module_symbol(db, self.canonical_module(), self.as_str())
+            .ignore_possibly_unbound()
+            .unwrap_or(Type::Unknown)
     }
 
     /// Return the module in which we should look up the definition for this class
@@ -2853,7 +2855,7 @@ impl<'db> TupleType<'db> {
 }
 
 #[cfg(test)]
-mod tests {
+pub(crate) mod tests {
     use super::*;
     use crate::db::tests::TestDb;
     use crate::program::{Program, SearchPathSettings};
@@ -2874,7 +2876,7 @@ mod tests {
         assert_eq!(size_of::<Type>(), 16);
     }
 
-    fn setup_db() -> TestDb {
+    pub(crate) fn setup_db() -> TestDb {
         let db = TestDb::new();
 
         let src_root = SystemPathBuf::from("/src");

crates/red_knot_python_semantic/src/types/infer.rs

@@ -1237,7 +1237,7 @@ impl<'db> TypeInferenceBuilder<'db> {
                 Type::Unknown
             }
             (Symbol::Type(enter_ty, enter_boundness), exit) => {
-                if enter_boundness == Boundness::MayBeUnbound {
+                if enter_boundness == Boundness::PossiblyUnbound {
                     self.diagnostics.add(
                         context_expression.into(),
                         "invalid-context-manager",
@@ -1276,7 +1276,7 @@ impl<'db> TypeInferenceBuilder<'db> {
                     Symbol::Type(exit_ty, exit_boundness) => {
                         // TODO: Use the `exit_ty` to determine if any raised exception is suppressed.
 
-                        if exit_boundness == Boundness::MayBeUnbound {
+                        if exit_boundness == Boundness::PossiblyUnbound {
                             self.diagnostics.add(
                                 context_expression.into(),
                                 "invalid-context-manager",
@@ -1340,7 +1340,8 @@ impl<'db> TypeInferenceBuilder<'db> {
             // TODO should infer `ExceptionGroup` if all caught exceptions
             // are subclasses of `Exception` --Alex
             builtins_symbol(self.db, "BaseExceptionGroup")
-                .unwrap_or_unknown()
+                .ignore_possibly_unbound()
+                .unwrap_or(Type::Unknown)
                 .to_instance(self.db)
         } else {
             // TODO: anything that's a consistent subtype of
@@ -1710,7 +1711,7 @@ impl<'db> TypeInferenceBuilder<'db> {
 
                     return match boundness {
                         Boundness::Bound => augmented_return_ty,
-                        Boundness::MayBeUnbound => {
+                        Boundness::PossiblyUnbound => {
                             let left_ty = target_type;
                             let right_ty = value_type;
 
@@ -2008,10 +2009,10 @@ impl<'db> TypeInferenceBuilder<'db> {
                     } = alias;
 
                     // For possibly-unbound names, just eliminate Unbound from the type; we
-                    // must be in a bound path. TODO diagnostic for maybe-unbound import?
+                    // must be in a bound path. TODO diagnostic for possibly-unbound import?
                     module_ty
                         .member(self.db, &ast::name::Name::new(&name.id))
-                        .as_type()
+                        .ignore_possibly_unbound()
                         .unwrap_or_else(|| {
                             self.diagnostics.add(
                                 AnyNodeRef::Alias(alias),
@@ -2186,7 +2187,8 @@ impl<'db> TypeInferenceBuilder<'db> {
                 .unwrap_or_else(|| KnownClass::Int.to_instance(self.db)),
             ast::Number::Float(_) => KnownClass::Float.to_instance(self.db),
             ast::Number::Complex { .. } => builtins_symbol(self.db, "complex")
-                .unwrap_or_unknown()
+                .ignore_possibly_unbound()
+                .unwrap_or(Type::Unknown)
                 .to_instance(self.db),
         }
     }
@@ -2265,7 +2267,9 @@ impl<'db> TypeInferenceBuilder<'db> {
         &mut self,
         _literal: &ast::ExprEllipsisLiteral,
     ) -> Type<'db> {
-        builtins_symbol(self.db, "Ellipsis").unwrap_or_unknown()
+        builtins_symbol(self.db, "Ellipsis")
+            .ignore_possibly_unbound()
+            .unwrap_or(Type::Unknown)
     }
 
     fn infer_tuple_expression(&mut self, tuple: &ast::ExprTuple) -> Type<'db> {
@@ -2688,21 +2692,21 @@ impl<'db> TypeInferenceBuilder<'db> {
             };
 
             // Fallback to builtins (without infinite recursion if we're already in builtins.)
-            if global_symbol.may_be_unbound()
+            if global_symbol.possibly_unbound()
                 && Some(self.scope()) != builtins_module_scope(self.db)
             {
-                let mut symbol = builtins_symbol(self.db, name);
-                if symbol.is_unbound() && name == "reveal_type" {
+                let mut builtins_symbol = builtins_symbol(self.db, name);
+                if builtins_symbol.is_unbound() && name == "reveal_type" {
                     self.diagnostics.add(
                         name_node.into(),
                         "undefined-reveal",
                         format_args!(
                             "`reveal_type` used without importing it; this is allowed for debugging convenience but will fail at runtime"),
                     );
-                    symbol = typing_extensions_symbol(self.db, name);
+                    builtins_symbol = typing_extensions_symbol(self.db, name);
                 }
 
-                global_symbol.replace_unbound_with(self.db, &symbol)
+                global_symbol.or_fall_back_to(self.db, &builtins_symbol)
             } else {
                 global_symbol
             }
@@ -2742,10 +2746,10 @@ impl<'db> TypeInferenceBuilder<'db> {
 
         let bindings_ty = bindings_ty(self.db, definitions);
 
-        if boundness == Boundness::MayBeUnbound {
+        if boundness == Boundness::PossiblyUnbound {
             match self.lookup_name(name) {
                 Symbol::Type(looked_up_ty, looked_up_boundness) => {
-                    if looked_up_boundness == Boundness::MayBeUnbound {
+                    if looked_up_boundness == Boundness::PossiblyUnbound {
                         self.diagnostics.add_possibly_unresolved_reference(name);
                     }
 
@@ -2787,7 +2791,8 @@ impl<'db> TypeInferenceBuilder<'db> {
         let value_ty = self.infer_expression(value);
         value_ty
             .member(self.db, &Name::new(&attr.id))
-            .unwrap_or_unknown()
+            .ignore_possibly_unbound()
+            .unwrap_or(Type::Unknown)
     }
 
     fn infer_attribute_expression(&mut self, attribute: &ast::ExprAttribute) -> Type<'db> {
@@ -3850,7 +3855,7 @@ impl<'db> TypeInferenceBuilder<'db> {
                 match value_meta_ty.member(self.db, "__getitem__") {
                     Symbol::Unbound => {}
                     Symbol::Type(dunder_getitem_method, boundness) => {
-                        if boundness == Boundness::MayBeUnbound {
+                        if boundness == Boundness::PossiblyUnbound {
                             self.diagnostics.add(
                                 value_node.into(),
                                 "call-possibly-unbound-method",
@@ -3894,7 +3899,7 @@ impl<'db> TypeInferenceBuilder<'db> {
                     match dunder_class_getitem_method {
                         Symbol::Unbound => {}
                         Symbol::Type(ty, boundness) => {
-                            if boundness == Boundness::MayBeUnbound {
+                            if boundness == Boundness::PossiblyUnbound {
                                 self.diagnostics.add(
                                     value_node.into(),
                                     "call-possibly-unbound-method",
@@ -4383,7 +4388,10 @@ impl<'db> TypeInferenceBuilder<'db> {
             ast::Expr::Attribute(ast::ExprAttribute { value, attr, .. }) => {
                 let value_ty = self.infer_expression(value);
                 // TODO: Check that value type is enum otherwise return None
-                value_ty.member(self.db, &attr.id).unwrap_or_unknown()
+                value_ty
+                    .member(self.db, &attr.id)
+                    .ignore_possibly_unbound()
+                    .unwrap_or(Type::Unknown)
             }
             ast::Expr::NoneLiteral(_) => Type::none(self.db),
             // for negative and positive numbers
@@ -4740,7 +4748,9 @@ mod tests {
             assert_eq!(scope.name(db), *expected_scope_name);
         }
 
-        let ty = symbol(db, scope, symbol_name).unwrap_or_unknown();
+        let ty = symbol(db, scope, symbol_name)
+            .ignore_possibly_unbound()
+            .unwrap_or(Type::Unknown);
         assert_eq!(ty.display(db).to_string(), expected);
     }