[ty] Narrow type context during collection literal inference (#23844)

Collection literals are effectively a special case of generic calls, so
there's no reason we shouldn't be applying the same narrowing logic
here.

Part of astral-sh/ty#3001.

Author: ibraheemdev

crates/ty_python_semantic/resources/mdtest/assignment/annotations.md

@@ -445,14 +445,21 @@ reveal_type(x8)  # revealed: Literal[True]
 x9: int | str = f2(True)
 reveal_type(x9)  # revealed: Literal[True]
 
-# TODO: Should not error. We could choose a concrete type here (pyright arbitrarily picks the
-# first), or keep the union (pyrefly does this). Mypy infers `list[int]` and errors.
-# error: [invalid-assignment]
 x10: list[int | str] | list[int | None] = [1, 2, 3]
-reveal_type(x10)  # revealed: list[int | str] | list[int | None]
+reveal_type(x10)  # revealed: list[int | str]
 
 x11: Sequence[int | str] | Sequence[int | None] = [1, 2, 3]
 reveal_type(x11)  # revealed: list[int]
+
+x12: list[int] | list[int | None] | list[str | None] = ["1", "2"]
+reveal_type(x12)  # revealed: list[str | None]
+
+x13: dict[str, list[int | None]] | dict[str, list[str | None]] = {"a": ["b"]}
+reveal_type(x13)  # revealed: dict[str, list[str | None]]
+
+x14 = [{"a": [1], "b": 1}, {"a": [1]}]
+x14.append(reveal_type({"b": 1}))  # revealed: dict[str, list[int] | int]
+reveal_type(x14)  # revealed: list[dict[str, list[int] | int] | dict[str, list[int]]]
 ```
 
 ## Annotations influence generic call argument inference

crates/ty_python_semantic/resources/mdtest/bidirectional.md

@@ -420,6 +420,26 @@ def _(x: int):
     y(lst(True), [z])
 ```
 
+```py
+def g[T](x: T, y: list[T | None]) -> T:
+    return x
+
+def _(flag: bool):
+    if flag:
+        x = 1
+
+    # error: [possibly-unresolved-reference]
+    x1: int | str = g(x, [1])
+    reveal_type(x1)  # revealed: int
+
+    if flag:
+        y = "1"
+
+    # error: [possibly-unresolved-reference]
+    x2: list[int | None] | list[str | None] = [y]
+    reveal_type(x2)  # revealed: list[str | None]
+```
+
 ```py
 class Bar(TypedDict):
     bar: int

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

@@ -4976,13 +4976,13 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> {
             _ => &[],
         };
 
-        // We silence diagnostics until we successfully narrow to a specific type.
-        let was_in_multi_inference = self.context.set_multi_inference(true);
-
         let mut try_narrow = |narrowed_ty| {
             let mut speculated_bindings = bindings.clone();
             let narrowed_tcx = TypeContext::new(Some(narrowed_ty));
 
+            // We silence diagnostics until we successfully narrow to a specific type.
+            let was_in_multi_inference = self.context.set_multi_inference(true);
+
             // Attempt to infer the argument types using the narrowed type context.
             self.infer_all_argument_types(
                 ast_arguments.clone(),
@@ -4993,6 +4993,9 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> {
                 MultiInferenceState::Ignore,
             );
 
+            // Restore the multi-inference state.
+            self.context.set_multi_inference(was_in_multi_inference);
+
             // Ensure the argument types match their annotated types.
             if speculated_bindings
                 .check_types_impl(
@@ -5023,8 +5026,6 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> {
             //
             // If necessary, infer the argument types again with diagnostics enabled.
             if !was_in_multi_inference {
-                self.context.set_multi_inference(was_in_multi_inference);
-
                 self.infer_all_argument_types(
                     ast_arguments.clone(),
                     argument_types,
@@ -5067,9 +5068,7 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> {
             }
         }
 
-        // Re-enable diagnostics, and infer against the entire union as a fallback.
-        self.context.set_multi_inference(was_in_multi_inference);
-
+        // Infer against the entire union as a fallback.
         self.infer_all_argument_types(
             ast_arguments,
             argument_types,
@@ -5784,11 +5783,11 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> {
             ctx: _,
         } = list;
 
-        let mut elts = elts.iter().map(|elt| [Some(elt)]);
+        let elts = elts.iter().map(|elt| [Some(elt)]).collect_vec();
         let mut infer_elt_ty =
             |builder: &mut Self, (_, elt, tcx)| builder.infer_expression(elt, tcx);
 
-        self.infer_collection_literal(KnownClass::List, &mut elts, &mut infer_elt_ty, tcx)
+        self.infer_collection_literal(KnownClass::List, &elts, &mut infer_elt_ty, tcx)
             .unwrap_or_else(|| {
                 KnownClass::List.to_specialized_instance(self.db(), &[Type::unknown()])
             })
@@ -5801,11 +5800,11 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> {
             elts,
         } = set;
 
-        let mut elts = elts.iter().map(|elt| [Some(elt)]);
+        let elts = elts.iter().map(|elt| [Some(elt)]).collect_vec();
         let mut infer_elt_ty =
             |builder: &mut Self, (_, elt, tcx)| builder.infer_expression(elt, tcx);
 
-        self.infer_collection_literal(KnownClass::Set, &mut elts, &mut infer_elt_ty, tcx)
+        self.infer_collection_literal(KnownClass::Set, &elts, &mut infer_elt_ty, tcx)
             .unwrap_or_else(|| {
                 KnownClass::Set.to_specialized_instance(self.db(), &[Type::unknown()])
             })
@@ -5888,9 +5887,10 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> {
                 .to_specialized_instance(self.db(), &[Type::unknown(), Type::unknown()]);
         }
 
-        let mut items = items
+        let items = items
             .iter()
-            .map(|item| [item.key.as_ref(), Some(&item.value)]);
+            .map(|item| [item.key.as_ref(), Some(&item.value)])
+            .collect_vec();
 
         // Avoid inferring the items multiple times if we already attempted to infer the
         // dictionary literal as a `TypedDict`. This also allows us to infer using the
@@ -5902,7 +5902,7 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> {
                 .unwrap_or_else(|| builder.infer_expression(elt, tcx))
         };
 
-        self.infer_collection_literal(KnownClass::Dict, &mut items, &mut infer_elt_ty, tcx)
+        self.infer_collection_literal(KnownClass::Dict, &items, &mut infer_elt_ty, tcx)
             .unwrap_or_else(|| {
                 KnownClass::Dict
                     .to_specialized_instance(self.db(), &[Type::unknown(), Type::unknown()])
@@ -5955,7 +5955,74 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> {
     fn infer_collection_literal<'expr, const N: usize>(
         &mut self,
         collection_class: KnownClass,
-        elts: &mut dyn Iterator<Item = [Option<&'expr ast::Expr>; N]>,
+        elts: &[[Option<&'expr ast::Expr>; N]],
+        infer_elt_expression: &mut dyn FnMut(&mut Self, ArgExpr<'db, 'expr>) -> Type<'db>,
+        tcx: TypeContext<'db>,
+    ) -> Option<Type<'db>> {
+        let db = self.db();
+
+        // If the type context is a union, attempt to narrow to a specific element.
+        let narrow_targets: &[_] = match tcx.annotation {
+            Some(Type::Union(union)) => union.elements(db),
+            _ => &[],
+        };
+
+        let mut try_narrow = |narrowed_ty| {
+            let narrowed_tcx = TypeContext::new(Some(narrowed_ty));
+
+            // We silence diagnostics until we successfully narrow to a specific type.
+            let prev_multi_inference = self.set_multi_inference_state(MultiInferenceState::Ignore);
+            let was_in_multi_inference = self.context.set_multi_inference(true);
+
+            // Attempt to infer the collection literal using the narrowed type context.
+            let inferred_ty = self.infer_collection_literal_impl(
+                collection_class,
+                elts,
+                infer_elt_expression,
+                narrowed_tcx,
+            )?;
+
+            // Restore the multi-inference state.
+            self.context.set_multi_inference(was_in_multi_inference);
+            self.set_multi_inference_state(prev_multi_inference);
+
+            // Ensure the inferred return type is assignable to the (narrowed) declared type.
+            if !inferred_ty.is_assignable_to(db, narrowed_ty) {
+                return None;
+            }
+
+            // Successfully narrowed to an element of the union.
+            //
+            // If necessary, infer the collection literal again with diagnostics enabled.
+            if !was_in_multi_inference {
+                self.infer_collection_literal_impl(
+                    collection_class,
+                    elts,
+                    infer_elt_expression,
+                    narrowed_tcx,
+                );
+            }
+
+            Some(inferred_ty)
+        };
+
+        for narrowed_ty in narrow_targets
+            .iter()
+            .filter(|ty| ty.class_specialization(db).is_some())
+        {
+            if let Some(result) = try_narrow(*narrowed_ty) {
+                return Some(result);
+            }
+        }
+
+        self.infer_collection_literal_impl(collection_class, elts, infer_elt_expression, tcx)
+    }
+
+    // Infer the type of a collection literal expression.
+    fn infer_collection_literal_impl<'expr, const N: usize>(
+        &mut self,
+        collection_class: KnownClass,
+        elts: &[[Option<&'expr ast::Expr>; N]],
         infer_elt_expression: &mut dyn FnMut(&mut Self, ArgExpr<'db, 'expr>) -> Type<'db>,
         tcx: TypeContext<'db>,
     ) -> Option<Type<'db>> {
@@ -5980,7 +6047,7 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> {
         let Some((collection_alias, generic_context, elt_tys)) = elt_tys(collection_class) else {
             // Infer the element types without type context, and fallback to `Unknown` for
             // custom typesheds.
-            for (i, elt) in elts.flatten().flatten().enumerate() {
+            for (i, elt) in elts.iter().flatten().flatten().enumerate() {
                 infer_elt_expression(self, (i, elt, TypeContext::default()));
             }
 
@@ -6009,12 +6076,7 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> {
                 FxHashMap::default();
 
             if let Some(tcx) = tcx.annotation
-                // If there are multiple potential type contexts, we fallback to `Unknown`.
-                // TODO: We could perform multi-inference here.
-                && tcx
-                    .filter_union(self.db(), |ty| ty.class_specialization(self.db()).is_some())
-                    .class_specialization(self.db())
-                    .is_some()
+                && tcx.class_specialization(self.db()).is_some()
             {
                 let collection_instance =
                     Type::instance(self.db(), ClassType::Generic(collection_alias));
@@ -6245,15 +6307,14 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> {
     fn infer_comprehension_specialization<const N: usize>(
         &mut self,
         collection_class: KnownClass,
-        elements: &[Option<&ast::Expr>; N],
+        elements: [Option<&ast::Expr>; N],
         inference: &ScopeInference<'db>,
         tcx: TypeContext<'db>,
     ) -> Option<Type<'db>> {
-        let mut elements = [elements].into_iter().copied();
         let mut infer_element_ty =
             |_builder: &mut Self, (_, elt, _)| inference.expression_type(elt);
 
-        self.infer_collection_literal(collection_class, &mut elements, &mut infer_element_ty, tcx)
+        self.infer_collection_literal(collection_class, &[elements], &mut infer_element_ty, tcx)
     }
 
     fn infer_list_comprehension_expression(
@@ -6280,7 +6341,7 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> {
         let inference = infer_scope_types(self.db(), scope, tcx);
         self.extend_scope(inference);
 
-        self.infer_comprehension_specialization(KnownClass::List, &[Some(elt)], inference, tcx)
+        self.infer_comprehension_specialization(KnownClass::List, [Some(elt)], inference, tcx)
             .unwrap_or_else(|| {
                 KnownClass::List.to_specialized_instance(self.db(), &[Type::unknown()])
             })
@@ -6310,7 +6371,7 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> {
         let inference = infer_scope_types(self.db(), scope, tcx);
         self.extend_scope(inference);
 
-        self.infer_comprehension_specialization(KnownClass::Set, &[Some(elt)], inference, tcx)
+        self.infer_comprehension_specialization(KnownClass::Set, [Some(elt)], inference, tcx)
             .unwrap_or_else(|| {
                 KnownClass::Set.to_specialized_instance(self.db(), &[Type::unknown()])
             })
@@ -6343,7 +6404,7 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> {
 
         self.infer_comprehension_specialization(
             KnownClass::Dict,
-            &[Some(key), Some(value)],
+            [Some(key), Some(value)],
             inference,
             tcx,
         )
@@ -6378,10 +6439,10 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> {
         } = listcomp;
 
         // Infer the element type using the outer type context.
-        let mut elts = [[Some(elt.as_ref())]].into_iter();
+        let elts = [[Some(elt.as_ref())]];
         let mut infer_elt_ty =
             |builder: &mut Self, (_, elt, tcx)| builder.infer_expression(elt, tcx);
-        self.infer_collection_literal(KnownClass::List, &mut elts, &mut infer_elt_ty, tcx);
+        self.infer_collection_literal(KnownClass::List, &elts, &mut infer_elt_ty, tcx);
 
         self.infer_comprehensions(generators);
     }
@@ -6399,10 +6460,10 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> {
         } = setcomp;
 
         // Infer the element type using the outer type context.
-        let mut elts = [[Some(elt.as_ref())]].into_iter();
+        let elts = [[Some(elt.as_ref())]];
         let mut infer_elt_ty =
             |builder: &mut Self, (_, elt, tcx)| builder.infer_expression(elt, tcx);
-        self.infer_collection_literal(KnownClass::Set, &mut elts, &mut infer_elt_ty, tcx);
+        self.infer_collection_literal(KnownClass::Set, &elts, &mut infer_elt_ty, tcx);
 
         self.infer_comprehensions(generators);
     }
@@ -6421,10 +6482,10 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> {
         } = dictcomp;
 
         // Infer the key and value types using the outer type context.
-        let mut elts = [[Some(key.as_ref()), Some(value.as_ref())]].into_iter();
+        let elts = [[Some(key.as_ref()), Some(value.as_ref())]];
         let mut infer_elt_ty =
             |builder: &mut Self, (_, elt, tcx)| builder.infer_expression(elt, tcx);
-        self.infer_collection_literal(KnownClass::Dict, &mut elts, &mut infer_elt_ty, tcx);
+        self.infer_collection_literal(KnownClass::Dict, &elts, &mut infer_elt_ty, tcx);
 
         self.infer_comprehensions(generators);
     }