[ty] Literal promotion refactor

crates/ty_python_semantic/resources/mdtest/protocols.md

@@ -893,8 +893,10 @@ class LotsOfBindings(Protocol):
     match object():
         case l:  # error: [ambiguous-protocol-member]
             ...
+    # error: [ambiguous-protocol-member] "Consider adding an annotation, e.g. `m: int | str = ...`"
+    m = 1 if 1.2 > 3.4 else "a"
 
-# revealed: frozenset[Literal["Nested", "NestedProtocol", "a", "b", "c", "d", "e", "f", "g", "h", "i", "j", "k", "l"]]
+# revealed: frozenset[Literal["Nested", "NestedProtocol", "a", "b", "c", "d", "e", "f", "g", "h", "i", "j", "k", "l", "m"]]
 reveal_type(get_protocol_members(LotsOfBindings))
 
 class Foo(Protocol):

crates/ty_python_semantic/src/types.rs

@@ -1166,21 +1166,26 @@ impl<'db> Type<'db> {
         }
     }
 
-    /// If this type is a literal, promote it to a type that this literal is an instance of.
+    /// Promote (possibly nested) literals to types that these literals are instances of.
     ///
     /// Note that this function tries to promote literals to a more user-friendly form than their
     /// fallback instance type. For example, `def _() -> int` is promoted to `Callable[[], int]`,
     /// as opposed to `FunctionType`.
-    pub(crate) fn literal_promotion_type(self, db: &'db dyn Db) -> Option<Type<'db>> {
+    pub(crate) fn promote_literals(self, db: &'db dyn Db) -> Type<'db> {
+        self.apply_type_mapping(db, &TypeMapping::PromoteLiterals)
+    }
+
+    /// Like [`Type::promote_literals`], but does not recurse into nested types.
+    fn promote_literals_impl(self, db: &'db dyn Db) -> Type<'db> {
         match self {
-            Type::StringLiteral(_) | Type::LiteralString => Some(KnownClass::Str.to_instance(db)),
-            Type::BooleanLiteral(_) => Some(KnownClass::Bool.to_instance(db)),
-            Type::IntLiteral(_) => Some(KnownClass::Int.to_instance(db)),
-            Type::BytesLiteral(_) => Some(KnownClass::Bytes.to_instance(db)),
-            Type::ModuleLiteral(_) => Some(KnownClass::ModuleType.to_instance(db)),
-            Type::EnumLiteral(literal) => Some(literal.enum_class_instance(db)),
-            Type::FunctionLiteral(literal) => Some(Type::Callable(literal.into_callable_type(db))),
-            _ => None,
+            Type::StringLiteral(_) | Type::LiteralString => KnownClass::Str.to_instance(db),
+            Type::BooleanLiteral(_) => KnownClass::Bool.to_instance(db),
+            Type::IntLiteral(_) => KnownClass::Int.to_instance(db),
+            Type::BytesLiteral(_) => KnownClass::Bytes.to_instance(db),
+            Type::ModuleLiteral(_) => KnownClass::ModuleType.to_instance(db),
+            Type::EnumLiteral(literal) => literal.enum_class_instance(db),
+            Type::FunctionLiteral(literal) => Type::Callable(literal.into_callable_type(db)),
+            _ => self,
         }
     }
 
@@ -6080,8 +6085,7 @@ impl<'db> Type<'db> {
                 let function = Type::FunctionLiteral(function.apply_type_mapping_impl(db, type_mapping, visitor));
 
                 match type_mapping {
-                    TypeMapping::PromoteLiterals => function.literal_promotion_type(db)
-                        .expect("function literal should have a promotion type"),
+                    TypeMapping::PromoteLiterals => function.promote_literals_impl(db),
                     _ => function
                 }
             }
@@ -6189,8 +6193,7 @@ impl<'db> Type<'db> {
                 TypeMapping::ReplaceSelf { .. } |
                 TypeMapping::MarkTypeVarsInferable(_) |
                 TypeMapping::Materialize(_) => self,
-                TypeMapping::PromoteLiterals => self.literal_promotion_type(db)
-                    .expect("literal type should have a promotion type"),
+                TypeMapping::PromoteLiterals => self.promote_literals_impl(db)
             }
 
             Type::Dynamic(_) => match type_mapping {

crates/ty_python_semantic/src/types/diagnostic.rs

@@ -2625,6 +2625,12 @@ pub(crate) fn report_undeclared_protocol_member(
                 SubclassOfInner::Dynamic(_) => return false,
             },
             Type::NominalInstance(instance) => instance.class(db),
+            Type::Union(union) => {
+                return union
+                    .elements(db)
+                    .iter()
+                    .all(|elem| should_give_hint(db, *elem));
+            }
             _ => return false,
         };
 
@@ -2656,9 +2662,7 @@ pub(crate) fn report_undeclared_protocol_member(
     if definition.kind(db).is_unannotated_assignment() {
         let binding_type = binding_type(db, definition);
 
-        let suggestion = binding_type
-            .literal_promotion_type(db)
-            .unwrap_or(binding_type);
+        let suggestion = binding_type.promote_literals(db);
 
         if should_give_hint(db, suggestion) {
             diagnostic.set_primary_message(format_args!(

crates/ty_python_semantic/src/types/infer/builder.rs

@@ -92,7 +92,7 @@ use crate::types::{
     DynamicType, IntersectionBuilder, IntersectionType, KnownClass, KnownInstanceType,
     MemberLookupPolicy, MetaclassCandidate, PEP695TypeAliasType, Parameter, ParameterForm,
     Parameters, SpecialFormType, SubclassOfType, TrackedConstraintSet, Truthiness, Type,
-    TypeAliasType, TypeAndQualifiers, TypeContext, TypeMapping, TypeQualifiers,
+    TypeAliasType, TypeAndQualifiers, TypeContext, TypeQualifiers,
     TypeVarBoundOrConstraintsEvaluation, TypeVarDefaultEvaluation, TypeVarInstance, TypeVarKind,
     UnionBuilder, UnionType, binding_type, todo_type,
 };
@@ -5432,8 +5432,7 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> {
 
                 // Convert any element literals to their promoted type form to avoid excessively large
                 // unions for large nested list literals, which the constraint solver struggles with.
-                let inferred_elt_ty =
-                    inferred_elt_ty.apply_type_mapping(self.db(), &TypeMapping::PromoteLiterals);
+                let inferred_elt_ty = inferred_elt_ty.promote_literals(self.db());
 
                 builder
                     .infer(Type::TypeVar(*elt_ty), inferred_elt_ty)