[ty] Support multi-inference through type aliases (#25245)

Resolves astral-sh/ty#3487.

Author: ibraheemdev

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

@@ -435,7 +435,7 @@ python-version = "3.12"
 ```
 
 ```py
-from typing import Literal, Sequence
+from typing import Literal, Mapping, Sequence
 
 def f[T](x: T) -> list[T]:
     return [x]
@@ -482,9 +482,20 @@ 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]]]
+x14: Mapping[str, list[int | None]] | Mapping[str, list[str | None]] = {"a": ["b"]}
+reveal_type(x14)  # revealed: dict[str, list[str | None]]
+
+x15 = [{"a": [1], "b": 1}, {"a": [1]}]
+x15.append(reveal_type({"b": 1}))  # revealed: dict[str, list[int] | int]
+reveal_type(x15)  # revealed: list[dict[str, list[int] | int] | dict[str, list[int]]]
+
+type EitherList = list[int | str] | list[int | None]
+
+x16: EitherList = [None, None]
+reveal_type(x16)  # revealed: list[int | None]
+
+x17: EitherList = ["1", "2", "3"]
+reveal_type(x17)  # revealed: list[int | str]
 ```
 
 ## Annotations influence generic call argument inference
@@ -855,7 +866,7 @@ x19: dict[int, dict[str, int]] = defaultdict(dict)
 reveal_type(x19)  # revealed: defaultdict[int, dict[str, int]]
 ```
 
-## Narrow generic unions
+## Narrow union declared type for generic calls
 
 ```toml
 [environment]
@@ -907,6 +918,37 @@ def _(target: TargetWithTD):
     reveal_type(target)  # revealed: (TD, /) -> None
 ```
 
+```py
+from typing import Mapping, Sequence
+
+x1: list[int | str] | list[int | None] = list((1, 2, 3))
+reveal_type(x1)  # revealed: list[int | str]
+
+x2: Sequence[int | str] | Sequence[int | None] = list((1, 2, 3))
+reveal_type(x2)  # revealed: list[int]
+
+x3: list[int] | list[int | None] | list[str | None] = list(("1", "2"))
+reveal_type(x3)  # revealed: list[str | None]
+
+x4: dict[str, list[int | None]] | dict[str, list[str | None]] = dict([("a", ["b"])])
+reveal_type(x4)  # revealed: dict[str, list[str | None]]
+
+x5: Mapping[str, list[int | None]] | Mapping[str, list[str | None]] = dict([("a", ["b"])])
+reveal_type(x5)  # revealed: dict[str, list[str | None]]
+
+x6 = [dict([("a", list((1,))), ("b", 1)]), dict([("a", list((1,)))])]
+x6.append(reveal_type(dict([("b", 1)])))  # revealed: dict[str, list[int] | int]
+reveal_type(x6)  # revealed: list[dict[str, list[int] | int] | dict[str, list[int]]]
+
+type EitherList = list[int | str] | list[int | None]
+
+x7: EitherList = list((None, None))
+reveal_type(x7)  # revealed: list[int | None]
+
+x8: EitherList = list(("1", "2", "3"))
+reveal_type(x8)  # revealed: list[int | str]
+```
+
 ## Prefer the inferred type of non-generic classes
 
 ```toml

crates/ty_python_semantic/src/types/infer.rs

@@ -590,8 +590,8 @@ impl<'db> TypeContext<'db> {
         Self { annotation }
     }
 
-    // If the type annotation is a specialized instance of the given `KnownClass`, returns the
-    // specialization.
+    /// If the type annotation is a specialized instance of the given `KnownClass`, returns the
+    /// specialization.
     fn known_specialization(
         &self,
         db: &'db dyn Db,
@@ -611,6 +611,16 @@ impl<'db> TypeContext<'db> {
         self.annotation
             .is_some_and(|ty| ty.is_typealias_special_form())
     }
+
+    /// If the type annotation is a union, returns the target elements that it can be narrowed to.
+    pub(crate) fn narrow_targets(&self, db: &'db dyn Db) -> Option<&[Type<'db>]> {
+        self.annotation
+            .and_then(|ty| ty.as_union_like(db))
+            // TODO: We could theoretically attempt to narrow to every element of
+            // the power set of this union. However, this leads to an exponential
+            // explosion of inference attempts, and is rarely needed in practice.
+            .map(|union| union.elements(db))
+    }
 }
 
 impl<'db> From<Type<'db>> for TypeContext<'db> {

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

@@ -5181,16 +5181,12 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> {
             .any(|overload| overload.signature.generic_context.is_some());
 
         // If the type context is a union, attempt to narrow to a specific element.
-        let narrow_targets: &[_] = match call_expression_tcx.annotation {
-            // TODO: We could theoretically attempt to narrow to every element of
-            // the power set of this union. However, this leads to an exponential
-            // explosion of inference attempts, and is rarely needed in practice.
-            //
+        let narrow_targets = call_expression_tcx
+            .narrow_targets(db)
             // We only need to attempt narrowing on generic calls, otherwise the type
             // context has no effect.
-            Some(Type::Union(union)) if has_generic_context => union.elements(db),
-            _ => &[],
-        };
+            .filter(|_| has_generic_context)
+            .unwrap_or_default();
 
         let mut try_narrow = |narrowed_ty: Type<'db>| {
             // Short-circuit if there is no overload with a matching return type.
@@ -6266,12 +6262,6 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> {
     ) -> 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 mut speculative_builder = self.speculate();
 
@@ -6293,8 +6283,11 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> {
             Some(inferred_ty)
         };
 
-        for narrowed_ty in narrow_targets
-            .iter()
+        // If the type context is a union, attempt to narrow to a specific element.
+        for narrowed_ty in tcx
+            .narrow_targets(db)
+            .into_iter()
+            .flatten()
             .filter(|ty| ty.class_specialization(db).is_some())
         {
             if let Some(result) = try_narrow(*narrowed_ty) {