[ty] Use declared attribute types as type context (#21143)

## Summary

For example:
```py
class X:
    x: list[int | str]

def _(x: X):
    x.x = [1]
```

Resolves https://github.com/astral-sh/ty/issues/1375.

crates/ty_python_semantic/resources/mdtest/bidirectional.md

@@ -185,12 +185,12 @@ Declared attribute types:
 
 ```py
 class E:
-    e: list[Literal[1]]
+    a: list[Literal[1]]
+    b: list[Literal[1]]
 
 def _(e: E):
-    # TODO: Implement attribute type context.
+    e.a = [1]
-    # error: [invalid-assignment] "Object of type `list[Unknown | int]` is not assignable to attribute `e` of type `list[Literal[1]]`"
+    E.b = [1]
-    e.e = [1]
 ```
 
 Function return types:
@@ -200,6 +200,41 @@ def f() -> list[Literal[1]]:
     return [1]
 ```
 
+## Instance attribute
+
+```toml
+[environment]
+python-version = "3.12"
+```
+
+Both meta and class/instance attribute annotations are used as type context:
+
+```py
+from typing import Literal, Any
+
+class DataDescriptor:
+    def __get__(self, instance: object, owner: type | None = None) -> list[Literal[1]]:
+        return []
+
+    def __set__(self, instance: object, value: list[Literal[1]]) -> None:
+        pass
+
+def lst[T](x: T) -> list[T]:
+    return [x]
+
+def _(flag: bool):
+    class Meta(type):
+        if flag:
+            x: DataDescriptor = DataDescriptor()
+
+    class C(metaclass=Meta):
+        x: list[int | None]
+
+    def _(c: C):
+        c.x = lst(1)
+        C.x = lst(1)
+```
+
 ## Class constructor parameters
 
 ```toml
@@ -226,3 +261,72 @@ A(f(1))
 # error: [invalid-argument-type] "Argument to bound method `__init__` is incorrect: Expected `list[int | None]`, found `list[list[Unknown]]`"
 A(f([]))
 ```
+
+## Multi-inference diagnostics
+
+```toml
+[environment]
+python-version = "3.12"
+```
+
+Diagnostics unrelated to the type-context are only reported once:
+
+`call.py`:
+
+```py
+def f[T](x: T) -> list[T]:
+    return [x]
+
+def a(x: list[bool], y: list[bool]): ...
+def b(x: list[int], y: list[int]): ...
+def c(x: list[int], y: list[int]): ...
+def _(x: int):
+    if x == 0:
+        y = a
+    elif x == 1:
+        y = b
+    else:
+        y = c
+
+    if x == 0:
+        z = True
+
+    y(f(True), [True])
+
+    # error: [possibly-unresolved-reference] "Name `z` used when possibly not defined"
+    y(f(True), [z])
+```
+
+`call_standalone_expression.py`:
+
+```py
+def f(_: str): ...
+def g(_: str): ...
+def _(a: object, b: object, flag: bool):
+    if flag:
+        x = f
+    else:
+        x = g
+
+    # error: [unsupported-operator] "Operator `>` is not supported for types `object` and `object`"
+    x(f"{'a' if a > b else 'b'}")
+```
+
+`attribute_assignment.py`:
+
+```py
+from typing import TypedDict
+
+class TD(TypedDict):
+    y: int
+
+class X:
+    td: TD
+
+def _(x: X, flag: bool):
+    if flag:
+        y = 1
+
+    # error: [possibly-unresolved-reference] "Name `y` used when possibly not defined"
+    x.td = {"y": y}
+```

crates/ty_python_semantic/resources/mdtest/call/union.md

@@ -281,46 +281,3 @@ def _(flag: bool):
     # we currently consider `TypedDict` instances to be subtypes of `dict`
     f({"y": 1})
 ```
-
-Diagnostics unrelated to the type-context are only reported once:
-
-`expression.py`:
-
-```py
-def f[T](x: T) -> list[T]:
-    return [x]
-
-def a(x: list[bool], y: list[bool]): ...
-def b(x: list[int], y: list[int]): ...
-def c(x: list[int], y: list[int]): ...
-def _(x: int):
-    if x == 0:
-        y = a
-    elif x == 1:
-        y = b
-    else:
-        y = c
-
-    if x == 0:
-        z = True
-
-    y(f(True), [True])
-
-    # error: [possibly-unresolved-reference] "Name `z` used when possibly not defined"
-    y(f(True), [z])
-```
-
-`standalone_expression.py`:
-
-```py
-def f(_: str): ...
-def g(_: str): ...
-def _(a: object, b: object, flag: bool):
-    if flag:
-        x = f
-    else:
-        x = g
-
-    # error: [unsupported-operator] "Operator `>` is not supported for types `object` and `object`"
-    x(f"{'a' if a > b else 'b'}")
-```

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

@@ -2924,12 +2924,12 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> {
         for item in items {
             let target = item.optional_vars.as_deref();
             if let Some(target) = target {
-                self.infer_target(target, &item.context_expr, |builder| {
+                self.infer_target(target, &item.context_expr, |builder, tcx| {
                     // TODO: `infer_with_statement_definition` reports a diagnostic if `ctx_manager_ty` isn't a context manager
                     //  but only if the target is a name. We should report a diagnostic here if the target isn't a name:
                     //  `with not_context_manager as a.x: ...
                     builder
-                        .infer_standalone_expression(&item.context_expr, TypeContext::default())
+                        .infer_standalone_expression(&item.context_expr, tcx)
                         .enter(builder.db())
                 });
             } else {
@@ -3393,8 +3393,8 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> {
         } = assignment;
 
         for target in targets {
-            self.infer_target(target, value, |builder| {
+            self.infer_target(target, value, |builder, tcx| {
-                builder.infer_standalone_expression(value, TypeContext::default())
+                builder.infer_standalone_expression(value, tcx)
             });
         }
     }
@@ -3410,13 +3410,19 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> {
     /// `target`.
     fn infer_target<F>(&mut self, target: &ast::Expr, value: &ast::Expr, infer_value_expr: F)
     where
-        F: Fn(&mut Self) -> Type<'db>,
+        F: Fn(&mut Self, TypeContext<'db>) -> Type<'db>,
     {
-        let assigned_ty = match target {
+        match target {
-            ast::Expr::Name(_) => None,
+            ast::Expr::Name(_) => {
-            _ => Some(infer_value_expr(self)),
+                self.infer_target_impl(target, value, None);
-        };
+            }
-        self.infer_target_impl(target, value, assigned_ty);
+
+            _ => self.infer_target_impl(
+                target,
+                value,
+                Some(&|builder, tcx| infer_value_expr(builder, tcx)),
+            ),
+        }
     }
 
     /// Make sure that the subscript assignment `obj[slice] = value` is valid.
@@ -3568,16 +3574,53 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> {
         target: &ast::ExprAttribute,
         object_ty: Type<'db>,
         attribute: &str,
-        value_ty: Type<'db>,
+        infer_value_ty: &dyn Fn(&mut Self, TypeContext<'db>) -> Type<'db>,
         emit_diagnostics: bool,
     ) -> bool {
         let db = self.db();
 
-        let ensure_assignable_to = |attr_ty| -> bool {
+        let mut first_tcx = None;
+
+        // A wrapper over `infer_value_ty` that allows inferring the value type multiple times
+        // during attribute resolution.
+        let pure_infer_value_ty = infer_value_ty;
+        let mut infer_value_ty = |builder: &mut Self, tcx: TypeContext<'db>| -> Type<'db> {
+            // Overwrite the previously inferred value, preferring later inferences, which are
+            // likely more precise. Note that we still ensure each inference is assignable to
+            // its declared type, so this mainly affects the IDE hover type.
+            let prev_multi_inference_state = mem::replace(
+                &mut builder.multi_inference_state,
+                MultiInferenceState::Overwrite,
+            );
+
+            // If we are inferring the argument multiple times, silence diagnostics to avoid duplicated warnings.
+            let was_in_multi_inference = if let Some(first_tcx) = first_tcx {
+                // The first time we infer an argument during multi-inference must be without type context,
+                // to avoid leaking diagnostics for bidirectional inference attempts.
+                debug_assert_eq!(first_tcx, TypeContext::default());
+
+                builder.context.set_multi_inference(true)
+            } else {
+                builder.context.is_in_multi_inference()
+            };
+
+            let value_ty = pure_infer_value_ty(builder, tcx);
+
+            // Reset the multi-inference state.
+            first_tcx.get_or_insert(tcx);
+            builder.multi_inference_state = prev_multi_inference_state;
+            builder.context.set_multi_inference(was_in_multi_inference);
+
+            value_ty
+        };
+
+        // This closure should only be called if `value_ty` was inferred with `attr_ty` as type context.
+        let ensure_assignable_to =
+            |builder: &Self, value_ty: Type<'db>, attr_ty: Type<'db>| -> bool {
                 let assignable = value_ty.is_assignable_to(db, attr_ty);
                 if !assignable && emit_diagnostics {
                     report_invalid_attribute_assignment(
-                    &self.context,
+                        &builder.context,
                         target.into(),
                         attr_ty,
                         value_ty,
@@ -3588,10 +3631,11 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> {
             };
 
         // Return true (and emit a diagnostic) if this is an invalid assignment to a `Final` attribute.
-        let invalid_assignment_to_final = |qualifiers: TypeQualifiers| -> bool {
+        let invalid_assignment_to_final = |builder: &Self, qualifiers: TypeQualifiers| -> bool {
             if qualifiers.contains(TypeQualifiers::FINAL) {
                 if emit_diagnostics {
-                    if let Some(builder) = self.context.report_lint(&INVALID_ASSIGNMENT, target) {
+                    if let Some(builder) = builder.context.report_lint(&INVALID_ASSIGNMENT, target)
+                    {
                         builder.into_diagnostic(format_args!(
                             "Cannot assign to final attribute `{attribute}` \
                                          on type `{}`",
@@ -3607,8 +3651,17 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> {
 
         match object_ty {
             Type::Union(union) => {
+                // TODO: We could perform multi-inference here with each element of the union as type context.
+                let value_ty = infer_value_ty(self, TypeContext::default());
+
                 if union.elements(self.db()).iter().all(|elem| {
-                    self.validate_attribute_assignment(target, *elem, attribute, value_ty, false)
+                    self.validate_attribute_assignment(
+                        target,
+                        *elem,
+                        attribute,
+                        &|_, _| value_ty,
+                        false,
+                    )
                 }) {
                     true
                 } else {
@@ -3631,9 +3684,18 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> {
             }
 
             Type::Intersection(intersection) => {
+                // TODO: We could perform multi-inference here with each element of the union as type context.
+                let value_ty = infer_value_ty(self, TypeContext::default());
+
                 // TODO: Handle negative intersection elements
                 if intersection.positive(db).iter().any(|elem| {
-                    self.validate_attribute_assignment(target, *elem, attribute, value_ty, false)
+                    self.validate_attribute_assignment(
+                        target,
+                        *elem,
+                        attribute,
+                        &|_, _| value_ty,
+                        false,
+                    )
                 }) {
                     true
                 } else {
@@ -3657,12 +3719,14 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> {
                 target,
                 alias.value_type(self.db()),
                 attribute,
-                value_ty,
+                pure_infer_value_ty,
                 emit_diagnostics,
             ),
 
             // Super instances do not allow attribute assignment
             Type::NominalInstance(instance) if instance.has_known_class(db, KnownClass::Super) => {
+                infer_value_ty(self, TypeContext::default());
+
                 if emit_diagnostics {
                     if let Some(builder) = self.context.report_lint(&INVALID_ASSIGNMENT, target) {
                         builder.into_diagnostic(format_args!(
@@ -3674,6 +3738,8 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> {
                 false
             }
             Type::BoundSuper(_) => {
+                infer_value_ty(self, TypeContext::default());
+
                 if emit_diagnostics {
                     if let Some(builder) = self.context.report_lint(&INVALID_ASSIGNMENT, target) {
                         builder.into_diagnostic(format_args!(
@@ -3685,7 +3751,10 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> {
                 false
             }
 
-            Type::Dynamic(..) | Type::Never => true,
+            Type::Dynamic(..) | Type::Never => {
+                infer_value_ty(self, TypeContext::default());
+                true
+            }
 
             Type::NominalInstance(..)
             | Type::ProtocolInstance(_)
@@ -3710,6 +3779,10 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> {
             | Type::AlwaysFalsy
             | Type::TypeIs(_)
             | Type::TypedDict(_) => {
+                // TODO: We could use the annotated parameter type of `__setattr__` as type context here.
+                // However, we would still have to perform the first inference without type context.
+                let value_ty = infer_value_ty(self, TypeContext::default());
+
                 // First, try to call the `__setattr__` dunder method. If this is present/defined, overrides
                 // assigning the attributed by the normal mechanism.
                 let setattr_dunder_call_result = object_ty.try_call_dunder_with_policy(
@@ -3811,7 +3884,7 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> {
                                 place: Place::Defined(meta_attr_ty, _, meta_attr_boundness),
                                 qualifiers,
                             } => {
-                                if invalid_assignment_to_final(qualifiers) {
+                                if invalid_assignment_to_final(self, qualifiers) {
                                     return false;
                                 }
 
@@ -3819,6 +3892,8 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> {
                                     if let Place::Defined(meta_dunder_set, _, _) =
                                         meta_attr_ty.class_member(db, "__set__".into()).place
                                     {
+                                        // TODO: We could use the annotated parameter type of `__set__` as
+                                        // type context here.
                                         let dunder_set_result = meta_dunder_set.try_call(
                                             db,
                                             &CallArguments::positional([
@@ -3844,7 +3919,12 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> {
 
                                         dunder_set_result.is_ok()
                                     } else {
-                                        ensure_assignable_to(meta_attr_ty)
+                                        let value_ty = infer_value_ty(
+                                            self,
+                                            TypeContext::new(Some(meta_attr_ty)),
+                                        );
+
+                                        ensure_assignable_to(self, value_ty, meta_attr_ty)
                                     };
 
                                 let assignable_to_instance_attribute = if meta_attr_boundness
@@ -3857,12 +3937,16 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> {
                                     } =
                                         object_ty.instance_member(db, attribute)
                                     {
-                                        if invalid_assignment_to_final(qualifiers) {
+                                        let value_ty = infer_value_ty(
+                                            self,
+                                            TypeContext::new(Some(instance_attr_ty)),
+                                        );
+                                        if invalid_assignment_to_final(self, qualifiers) {
                                             return false;
                                         }
 
                                         (
-                                            ensure_assignable_to(instance_attr_ty),
+                                            ensure_assignable_to(self, value_ty, instance_attr_ty),
                                             instance_attr_boundness,
                                         )
                                     } else {
@@ -3896,7 +3980,11 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> {
                                     qualifiers,
                                 } = object_ty.instance_member(db, attribute)
                                 {
-                                    if invalid_assignment_to_final(qualifiers) {
+                                    let value_ty = infer_value_ty(
+                                        self,
+                                        TypeContext::new(Some(instance_attr_ty)),
+                                    );
+                                    if invalid_assignment_to_final(self, qualifiers) {
                                         return false;
                                     }
 
@@ -3909,7 +3997,7 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> {
                                         );
                                     }
 
-                                    ensure_assignable_to(instance_attr_ty)
+                                    ensure_assignable_to(self, value_ty, instance_attr_ty)
                                 } else {
                                     if emit_diagnostics {
                                         if let Some(builder) =
@@ -3937,13 +4025,19 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> {
                         place: Place::Defined(meta_attr_ty, _, meta_attr_boundness),
                         qualifiers,
                     } => {
-                        if invalid_assignment_to_final(qualifiers) {
+                        // We may have to perform multi-inference if the meta attribute is possibly unbound.
+                        // However, we are required to perform the first inference without type context.
+                        let value_ty = infer_value_ty(self, TypeContext::default());
+
+                        if invalid_assignment_to_final(self, qualifiers) {
                             return false;
                         }
 
                         let assignable_to_meta_attr = if let Place::Defined(meta_dunder_set, _, _) =
                             meta_attr_ty.class_member(db, "__set__".into()).place
                         {
+                            // TODO: We could use the annotated parameter type of `__set__` as
+                            // type context here.
                             let dunder_set_result = meta_dunder_set.try_call(
                                 db,
                                 &CallArguments::positional([meta_attr_ty, object_ty, value_ty]),
@@ -3963,7 +4057,9 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> {
 
                             dunder_set_result.is_ok()
                         } else {
-                            ensure_assignable_to(meta_attr_ty)
+                            let value_ty =
+                                infer_value_ty(self, TypeContext::new(Some(meta_attr_ty)));
+                            ensure_assignable_to(self, value_ty, meta_attr_ty)
                         };
 
                         let assignable_to_class_attr = if meta_attr_boundness
@@ -3976,7 +4072,12 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> {
                                         .expect("called on Type::ClassLiteral or Type::SubclassOf")
                                         .place
                                 {
-                                    (ensure_assignable_to(class_attr_ty), class_attr_boundness)
+                                    let value_ty =
+                                        infer_value_ty(self, TypeContext::new(Some(class_attr_ty)));
+                                    (
+                                        ensure_assignable_to(self, value_ty, class_attr_ty),
+                                        class_attr_boundness,
+                                    )
                                 } else {
                                     (true, Definedness::PossiblyUndefined)
                                 };
@@ -4008,7 +4109,9 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> {
                             .find_name_in_mro(db, attribute)
                             .expect("called on Type::ClassLiteral or Type::SubclassOf")
                         {
-                            if invalid_assignment_to_final(qualifiers) {
+                            let value_ty =
+                                infer_value_ty(self, TypeContext::new(Some(class_attr_ty)));
+                            if invalid_assignment_to_final(self, qualifiers) {
                                 return false;
                             }
 
@@ -4021,8 +4124,10 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> {
                                 );
                             }
 
-                            ensure_assignable_to(class_attr_ty)
+                            ensure_assignable_to(self, value_ty, class_attr_ty)
                         } else {
+                            infer_value_ty(self, TypeContext::default());
+
                             let attribute_is_bound_on_instance =
                                 object_ty.to_instance(self.db()).is_some_and(|instance| {
                                     !instance
@@ -4064,6 +4169,8 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> {
 
             Type::ModuleLiteral(module) => {
                 if let Place::Defined(attr_ty, _, _) = module.static_member(db, attribute).place {
+                    let value_ty = infer_value_ty(self, TypeContext::new(Some(attr_ty)));
+
                     let assignable = value_ty.is_assignable_to(db, attr_ty);
                     if assignable {
                         true
@@ -4080,6 +4187,8 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> {
                         false
                     }
                 } else {
+                    infer_value_ty(self, TypeContext::default());
+
                     if emit_diagnostics {
                         if let Some(builder) =
                             self.context.report_lint(&UNRESOLVED_ATTRIBUTE, target)
@@ -4098,22 +4207,35 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> {
         }
     }
 
+    #[expect(clippy::type_complexity)]
     fn infer_target_impl(
         &mut self,
         target: &ast::Expr,
         value: &ast::Expr,
-        assigned_ty: Option<Type<'db>>,
+        infer_assigned_ty: Option<&dyn Fn(&mut Self, TypeContext<'db>) -> Type<'db>>,
     ) {
         match target {
-            ast::Expr::Name(name) => self.infer_definition(name),
+            ast::Expr::Name(name) => {
+                if let Some(infer_assigned_ty) = infer_assigned_ty {
+                    infer_assigned_ty(self, TypeContext::default());
+                }
+
+                self.infer_definition(name);
+            }
             ast::Expr::List(ast::ExprList { elts, .. })
             | ast::Expr::Tuple(ast::ExprTuple { elts, .. }) => {
+                let assigned_ty = infer_assigned_ty.map(|f| f(self, TypeContext::default()));
+
                 if let Some(tuple_spec) =
                     assigned_ty.and_then(|ty| ty.tuple_instance_spec(self.db()))
                 {
-                    let mut assigned_tys = tuple_spec.all_elements();
+                    let assigned_tys = tuple_spec.all_elements().copied().collect::<Vec<_>>();
-                    for element in elts {
+
-                        self.infer_target_impl(element, value, assigned_tys.next().copied());
+                    for (i, element) in elts.iter().enumerate() {
+                        match assigned_tys.get(i).copied() {
+                            None => self.infer_target_impl(element, value, None),
+                            Some(ty) => self.infer_target_impl(element, value, Some(&|_, _| ty)),
+                        }
                     }
                 } else {
                     for element in elts {
@@ -4129,29 +4251,39 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> {
                     ..
                 },
             ) => {
-                self.store_expression_type(target, assigned_ty.unwrap_or(Type::unknown()));
-
                 let object_ty = self.infer_expression(object, TypeContext::default());
 
-                if let Some(assigned_ty) = assigned_ty {
+                if let Some(infer_assigned_ty) = infer_assigned_ty {
+                    let infer_assigned_ty = &|builder: &mut Self, tcx| {
+                        let assigned_ty = infer_assigned_ty(builder, tcx);
+                        builder.store_expression_type(target, assigned_ty);
+                        assigned_ty
+                    };
+
                     self.validate_attribute_assignment(
                         attr_expr,
                         object_ty,
                         attr.id(),
-                        assigned_ty,
+                        infer_assigned_ty,
                         true,
                     );
                 }
             }
             ast::Expr::Subscript(subscript_expr) => {
+                let assigned_ty = infer_assigned_ty.map(|f| f(self, TypeContext::default()));
                 self.store_expression_type(target, assigned_ty.unwrap_or(Type::unknown()));
 
                 if let Some(assigned_ty) = assigned_ty {
                     self.validate_subscript_assignment(subscript_expr, value, assigned_ty);
                 }
             }
-            _ => {
+
             // TODO: Remove this once we handle all possible assignment targets.
+            _ => {
+                if let Some(infer_assigned_ty) = infer_assigned_ty {
+                    infer_assigned_ty(self, TypeContext::default());
+                }
+
                 self.infer_expression(target, TypeContext::default());
             }
         }
@@ -4836,12 +4968,12 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> {
             is_async: _,
         } = for_statement;
 
-        self.infer_target(target, iter, |builder| {
+        self.infer_target(target, iter, |builder, tcx| {
             // TODO: `infer_for_statement_definition` reports a diagnostic if `iter_ty` isn't iterable
             //  but only if the target is a name. We should report a diagnostic here if the target isn't a name:
             //  `for a.x in not_iterable: ...
             builder
-                .infer_standalone_expression(iter, TypeContext::default())
+                .infer_standalone_expression(iter, tcx)
                 .iterate(builder.db())
                 .homogeneous_element_type(builder.db())
         });
@@ -5863,6 +5995,10 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> {
                 assert_eq!(previous, None);
             }
 
+            MultiInferenceState::Overwrite => {
+                self.expressions.insert(expression.into(), ty);
+            }
+
             MultiInferenceState::Intersect => {
                 self.expressions
                     .entry(expression.into())
@@ -6430,7 +6566,7 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> {
             is_async: _,
         } = comprehension;
 
-        self.infer_target(target, iter, |builder| {
+        self.infer_target(target, iter, |builder, tcx| {
             // TODO: `infer_comprehension_definition` reports a diagnostic if `iter_ty` isn't iterable
             //  but only if the target is a name. We should report a diagnostic here if the target isn't a name:
             //  `[... for a.x in not_iterable]
@@ -6438,11 +6574,11 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> {
                 infer_same_file_expression_type(
                     builder.db(),
                     builder.index.expression(iter),
-                    TypeContext::default(),
+                    tcx,
                     builder.module(),
                 )
             } else {
-                builder.infer_standalone_expression(iter, TypeContext::default())
+                builder.infer_standalone_expression(iter, tcx)
             }
             .iterate(builder.db())
             .homogeneous_element_type(builder.db())
@@ -10153,16 +10289,16 @@ enum MultiInferenceState {
     #[default]
     Panic,
 
+    /// Overwrite the previously inferred value.
+    Overwrite,
+
     /// Store the intersection of all types inferred for the expression.
     Intersect,
 }
 
 impl MultiInferenceState {
-    fn is_panic(self) -> bool {
+    const fn is_panic(self) -> bool {
-        match self {
+        matches!(self, MultiInferenceState::Panic)
-            MultiInferenceState::Panic => true,
-            MultiInferenceState::Intersect => false,
-        }
     }
 }