[ty] Avoid widening compatible generic return alternatives (astral-sh#29578)

Since astral-sh#28303, when a generic call inference has multiple valid
solutions, we will intersect the return types resulting from
specializing with each valid solution. This approach has various caveats
discussed in the description of that PR. One of them is that sometimes a
call with an invariant return type (e.g. `list[T]`) can be validly
inferred as returning e.g. either `list[int]` or `list[str]`, but not
both simultaneously (consider if it actually returns `[]` -- it can be
viewed as either kind of list, but not both at the same time).
Intersecting `list[int]` and `list[str]` would result in `Never`, which
is wrong: the call can return.

For this reason, we skip intersecting return types if they vary in an
invariant type.

But in some cases one specialized return type is a subtype of all the
other ones. In this case it's safe to go ahead and prefer the narrower
type, even if they vary in invariant position and thus we wouldn't
otherwise intersect them.

This avoids a regression in astral-sh#29236.

Author: carljm

crates/ty_python_semantic/resources/mdtest/generics/legacy/functions.md

@@ -3793,7 +3793,8 @@ make independently valid specializations incompatible: a fresh `[]` can be typed
 `Never`, which would incorrectly imply that allocating the list cannot return.
 
 We handle this conservatively: if an invariant component differs across valid specializations, we do
-not intersect their return types.
+not intersect their return types. (But if one valid specialization results in a return type that is
+a subtype of all others, we select that one.)
 
 ### Mutable returns from intersection arguments
 
@@ -4000,6 +4001,36 @@ def _(
     reveal_type(invoke_all(reordered, All()))  # revealed: tuple[B, list[int]] | tuple[A, list[int]] | int
 ```
 
+The call to `invoke_all` below can infer three different types for `R`, one from each callback
+signature. Each choice gives a possible return type for the call. Inference selects `list[B]`
+because it is a subtype of both other return types, `list[A] | list[B]` and `list[B] | list[C]`:
+
+```py
+def _(
+    callback: Intersection[
+        Callable[[A], list[A] | list[B]],
+        Callable[[B], list[B] | list[C]],
+        Callable[[C], list[B]],
+    ],
+) -> None:
+    reveal_type(invoke_all(callback, All()))  # revealed: list[B]
+```
+
+Being a subtype of just one other return type is insufficient. Here `list[A]` and `list[B]` are both
+subtypes of `list[A] | list[B]`, but neither is a subtype of the other. Inference retains
+`list[A] | list[B]`:
+
+```py
+def _(
+    callback: Intersection[
+        Callable[[A], list[A] | list[B]],
+        Callable[[B], list[A]],
+        Callable[[C], list[B]],
+    ],
+) -> None:
+    reveal_type(invoke_all(callback, All()))  # revealed: list[A] | list[B]
+```
+
 For generic classes with both covariant and invariant parameters, refinement depends on which
 parameters vary. `Wrapper` only produces `T`, while its writable `value` attribute makes `U`
 invariant:
@@ -4029,6 +4060,55 @@ def _(
     reveal_type(invoke(varying, object()))  # revealed: Wrapper[A, int] | Wrapper[B, str]
 ```
 
+### Returning values from mixed-variance unions
+
+Inference finds two valid specializations for the call to `first_value` below, returning `V@forward`
+and `V@forward | Invariant[V@forward]`. The first return type is a subtype of the second, so
+inference selects `V@forward`:
+
+```py
+from typing import Callable, Generic, TypeVar
+
+V = TypeVar("V")
+V_co = TypeVar("V_co", covariant=True)
+
+class Covariant(Generic[V_co]):
+    def get(self) -> V_co:
+        raise NotImplementedError
+
+class Invariant(Generic[V]):
+    value: V
+
+def first_value(mapping: Callable[[], V | Invariant[V] | Covariant[V]]) -> V:
+    raise NotImplementedError
+
+def forward(mapping: V | Invariant[V] | Covariant[V]) -> V:
+    value = first_value(lambda: mapping)
+    reveal_type(value)  # revealed: V@forward
+    return value
+```
+
+The comparison uses the complete return types after substituting the inferred type arguments. For
+`make_list`, the two choices produce `list[V]` and `list[V | Invariant[V]]`. Neither is a subtype of
+the other, because lists are invariant, so inference keeps the merged result below.
+
+For `make_sink`, a callable accepting `V | Invariant[V]` is a subtype of one accepting only `V`: it
+accepts every input the latter accepts. The wider parameter type therefore produces the narrower
+callable type, which inference selects:
+
+```py
+def make_list(mapping: Callable[[], V | Invariant[V] | Covariant[V]]) -> list[V]:
+    return []
+
+def make_sink(mapping: Callable[[], V | Invariant[V] | Covariant[V]]) -> Callable[[V], None]:
+    return lambda value: None
+
+def wrapped_returns(mapping: V | Invariant[V] | Covariant[V]) -> None:
+    # revealed: list[V@wrapped_returns | Covariant[V@wrapped_returns] | Invariant[V@wrapped_returns]]
+    reveal_type(make_list(lambda: mapping))
+    reveal_type(make_sink(lambda: mapping))  # revealed: (Invariant[V@wrapped_returns] | V@wrapped_returns, /) -> None
+```
+
 ### Inferred TypedDict return alternatives
 
 Writable `TypedDict` fields are invariant even when the dictionaries have no type parameters. A

crates/ty_python_semantic/resources/mdtest/generics/pep695/functions.md

@@ -1354,7 +1354,8 @@ make independently valid specializations incompatible: a fresh `[]` can be typed
 `Never`, which would incorrectly imply that allocating the list cannot return.
 
 We handle this conservatively: if an invariant component differs across valid specializations, we do
-not intersect their return types.
+not intersect their return types. (But if one valid specialization results in a return type that is
+a subtype of all others, we select that one.)
 
 ### Mutable returns from intersection arguments
 
@@ -1540,6 +1541,36 @@ def _(
     reveal_type(invoke_all(reordered, All()))  # revealed: tuple[B, list[int]] | tuple[A, list[int]] | int
 ```
 
+The call to `invoke_all` below can infer three different types for `R`, one from each callback
+signature. Each choice gives a possible return type for the call. Inference selects `list[B]`
+because it is a subtype of both other return types, `list[A] | list[B]` and `list[B] | list[C]`:
+
+```py
+def _(
+    callback: Intersection[
+        Callable[[A], list[A] | list[B]],
+        Callable[[B], list[B] | list[C]],
+        Callable[[C], list[B]],
+    ],
+) -> None:
+    reveal_type(invoke_all(callback, All()))  # revealed: list[B]
+```
+
+Being a subtype of just one other return type is insufficient. Here `list[A]` and `list[B]` are both
+subtypes of `list[A] | list[B]`, but neither is a subtype of the other. Inference retains
+`list[A] | list[B]`:
+
+```py
+def _(
+    callback: Intersection[
+        Callable[[A], list[A] | list[B]],
+        Callable[[B], list[A]],
+        Callable[[C], list[B]],
+    ],
+) -> None:
+    reveal_type(invoke_all(callback, All()))  # revealed: list[A] | list[B]
+```
+
 For generic classes with both covariant and invariant parameters, refinement depends on which
 parameters vary. `Wrapper` only produces `T`, while its writable `value` attribute makes `U`
 invariant:
@@ -1564,6 +1595,52 @@ def _(
     reveal_type(invoke(varying, object()))  # revealed: Wrapper[A, int] | Wrapper[B, str]
 ```
 
+### Returning values from mixed-variance unions
+
+Inference finds two valid specializations for the call to `first_value` below, returning `V@forward`
+and `V@forward | Invariant[V@forward]`. The first return type is a subtype of the second, so
+inference selects `V@forward`:
+
+```py
+from typing import Callable
+
+class Covariant[V]:
+    def get(self) -> V:
+        raise NotImplementedError
+
+class Invariant[V]:
+    value: V
+
+def first_value[V](mapping: Callable[[], V | Invariant[V] | Covariant[V]]) -> V:
+    raise NotImplementedError
+
+def forward[V](mapping: V | Invariant[V] | Covariant[V]) -> V:
+    value = first_value(lambda: mapping)
+    reveal_type(value)  # revealed: V@forward
+    return value
+```
+
+The comparison uses the complete return types after substituting the inferred type arguments. For
+`make_list`, the two choices produce `list[V]` and `list[V | Invariant[V]]`. Neither is a subtype of
+the other, because lists are invariant, so inference keeps the merged result below.
+
+For `make_sink`, a callable accepting `V | Invariant[V]` is a subtype of one accepting only `V`: it
+accepts every input the latter accepts. The wider parameter type therefore produces the narrower
+callable type, which inference selects:
+
+```py
+def make_list[V](mapping: Callable[[], V | Invariant[V] | Covariant[V]]) -> list[V]:
+    return []
+
+def make_sink[V](mapping: Callable[[], V | Invariant[V] | Covariant[V]]) -> Callable[[V], None]:
+    return lambda value: None
+
+def wrapped_returns[V](mapping: V | Invariant[V] | Covariant[V]) -> None:
+    # revealed: list[V@wrapped_returns | Covariant[V@wrapped_returns] | Invariant[V@wrapped_returns]]
+    reveal_type(make_list(lambda: mapping))
+    reveal_type(make_sink(lambda: mapping))  # revealed: (Invariant[V@wrapped_returns] | V@wrapped_returns, /) -> None
+```
+
 ### Inferred TypedDict return alternatives
 
 Writable `TypedDict` fields are invariant even when the dictionaries have no type parameters. A

crates/ty_python_semantic/src/types/call/bind.rs

@@ -6384,6 +6384,39 @@ impl<'db> CallInference<'_, 'db> {
         let db = self.db;
         let env = self.env;
 
+        let mut returns = Vec::with_capacity(specializations.len());
+        let mut budget = ProjectionTypeBudget::new(SolutionBudget::default().type_terms);
+        for &specialization in specializations {
+            let return_ty = self.return_ty.apply_specialization(db, specialization);
+            budget.charge_type(db, return_ty).ok()?;
+            // Expose aliased unions to the intersection constructor's expansion budget.
+            returns.push(match return_ty.resolve_type_alias(db) {
+                Type::Union(union) => union.expand_aliases(db, env),
+                ty => ty,
+            });
+        }
+
+        // If one valid specialization's return type is a subtype of all the others, use it
+        // directly. For example, `int` can be selected over `int | list[int]`, even though the
+        // invariant-type check below would otherwise prevent intersecting them. Compare complete
+        // return types: substituting a narrower type argument can produce an incomparable
+        // (`list[T]`) or wider (`Callable[[T], None]`) return type.
+        //
+        // Find a candidate in one pass, then check it against every return type; two
+        // incomparable returns may still have a later return that is a subtype of both.
+        let mut candidate = *returns.first()?;
+        for &return_ty in returns.iter().skip(1) {
+            if return_ty != candidate && return_ty.is_subtype_of(db, env, candidate) {
+                candidate = return_ty;
+            }
+        }
+        if returns
+            .iter()
+            .all(|&return_ty| candidate == return_ty || candidate.is_subtype_of(db, env, return_ty))
+        {
+            return Some(candidate);
+        }
+
         // A fresh invariant container can receive its specialization from context: the same
         // expression `[]` is valid as either `list[A]` or `list[B]`. These typings make different,
         // mutually exclusive choices for the allocated list's static element type, rather than
@@ -6423,18 +6456,6 @@ impl<'db> CallInference<'_, 'db> {
                 })
         });
 
-        let mut returns = Vec::with_capacity(specializations.len());
-        let mut budget = ProjectionTypeBudget::new(SolutionBudget::default().type_terms);
-        for &specialization in specializations {
-            let return_ty = self.return_ty.apply_specialization(db, specialization);
-            budget.charge_type(db, return_ty).ok()?;
-            // Expose aliased unions to the intersection constructor's expansion budget.
-            returns.push(match return_ty.resolve_type_alias(db) {
-                Type::Union(union) => union.expand_aliases(db, env),
-                ty => ty,
-            });
-        }
-
         if has_invariant_difference {
             let merged = inference.merged_specialization(db);
             // Preserve the existing merged return when that specialization accepts the call.