Ignore return-only vars when an overload decides whether it wins

pyrefly/lib/alt/callable.rs

@@ -2553,6 +2553,11 @@ impl<'ctx, 'answer, Ans: LookupAnswer> AnswersSolver<'ctx, 'answer, Ans> {
         if let Some(self_qs) = self_qs {
             call_boundary.defer_quantified(self_qs);
         }
+        // A violation belongs to this call, so that an overload that causes one is rejected. It
+        // must be taken before generalizing, which expands the answer's var away.
+        let restriction_errors = ctor_targs.as_ref().map_or_else(Vec::new, |targs| {
+            self.solver().take_restricted_answer_errors(targs)
+        });
         if let Some(targs) = ctor_targs {
             let recorded_vars = call_context.captured_vars();
             self.solver().generalize_class_targs(targs, &recorded_vars);
@@ -2562,7 +2567,8 @@ impl<'ctx, 'answer, Ans: LookupAnswer> AnswersSolver<'ctx, 'answer, Ans> {
             self.type_order(),
             call_boundary,
         );
-        let errors = finish_result.map_or_else(|e| e.to_vec(), |_| Vec::new());
+        let mut errors = finish_result.map_or_else(|e| e.to_vec(), |_| Vec::new());
+        errors.extend(restriction_errors);
 
         // Apply meta-shape inference if bound args were collected
         let ret = if let Some(meta_shape_func) = meta_shape_func

pyrefly/lib/solver/solver.rs

@@ -151,6 +151,10 @@ impl Bounds {
 pub struct OverloadBranch {
     branch_index: usize,
     values: SmallMap<Var, Variable>,
+    /// The captured vars that appear in the parameters the branch was matched against. The other
+    /// captured vars only receive the branch's return type, so they cannot decide which branch
+    /// applies.
+    param_vars: SmallSet<Var>,
     /// Vars that may become free quantifieds, as of snapshot time. Read by
     /// `overload_branch_value_type` to decide whether an unsolved branch value
     /// should be a free quantified.
@@ -305,9 +309,11 @@ enum Variable {
 /// cannot enforce the parameter's restriction on its own. Argument matching checks each argument
 /// against `param` instead and records the first violation in `error`.
 ///
-/// `finish_quantified_with_captures` reports the violation and takes the whole record off the answer,
-/// which also keeps `param` out of published answers: `sanitize_vars` traverses only the answer
-/// type, so a parameter left here could hide a variable that never gets pinned.
+/// Each call that matches arguments against the answer takes the violation it caused, so that an
+/// overload that causes one is rejected like it is without the expected type.
+/// `finish_quantified_with_captures` reports any violation left and takes the whole record off the
+/// answer, which also keeps `param` out of published answers: `sanitize_vars` traverses only the
+/// answer type, so a parameter left here could hide a variable that never gets pinned.
 ///
 /// The violation is recorded here rather than in `Solver::instantiation_errors` because an entry in
 /// that map means "solving this variable went wrong", which makes `with_snapshot` reject the
@@ -1032,6 +1038,7 @@ impl Solver {
         &self,
         branch_index: usize,
         vars: &[Var],
+        param_vars: &SmallSet<Var>,
         free_quantified_vars_in_call: &SmallSet<Var>,
     ) -> OverloadBranch {
         let variables = self.variables.lock();
@@ -1052,6 +1059,7 @@ impl Solver {
         OverloadBranch {
             branch_index,
             values,
+            param_vars: param_vars.clone(),
             free_quantified_vars,
             errors,
         }
@@ -2062,9 +2070,11 @@ impl Solver {
             surviving_branches.push((capture.branch_index, state));
             // Step 5 of overload call evaluation: once the types passed to a branch's parameters
             // are known, a branch that accepts every materialization of them eliminates later
-            // branches. A var receives a parameter type as an upper bound, or is pinned outright.
+            // branches. A var in the parameters receives a parameter type as an upper bound, or is
+            // pinned outright.
             let passed_types_are_known = capture.values.iter().all(|(var, branch_value)| {
-                solved_vars.contains_key(var)
+                !capture.param_vars.contains(var)
+                    || solved_vars.contains_key(var)
                     || match branch_value {
                         Variable::Quantified { bounds, .. } | Variable::Unwrap(bounds) => {
                             bounds.upper.is_empty()
@@ -2463,6 +2473,29 @@ impl Solver {
         self.generalize_class_targs_impl(targs, vars_with_overload_branches, true)
     }
 
+    /// Take the restriction violations that argument matching recorded on the restricted answers
+    /// in `targs`. See [`RestrictedAnswer`].
+    ///
+    /// `targs` must not have been expanded yet, because expansion replaces each answer's var with
+    /// its type and leaves nothing to take from.
+    pub fn take_restricted_answer_errors(&self, targs: &TArgs) -> Vec<TypeVarSpecializationError> {
+        let lock = self.variables.lock();
+        targs
+            .as_slice()
+            .iter()
+            .filter_map(|t| {
+                let Type::Var(v) = t else { return None };
+                match &mut *lock.get_mut(*v) {
+                    Variable::Answer {
+                        restricted: Some(restricted),
+                        ..
+                    } => restricted.error.take(),
+                    _ => None,
+                }
+            })
+            .collect()
+    }
+
     fn generalize_class_targs_impl(
         &self,
         targs: &mut TArgs,
@@ -4176,8 +4209,8 @@ impl<'solver, 'subset, Ans: LookupAnswer> Subset<'solver, 'subset, Ans> {
                         ty: t2, restricted, ..
                     } => {
                         let t2 = t2.clone();
-                        // Only the first violation is kept, so a parameter that has already been
-                        // rejected needs no further checking.
+                        // Only the first violation per call is kept, so a parameter that this call
+                        // has already rejected needs no further checking.
                         let param = restricted
                             .as_ref()
                             .filter(|r| r.error.is_none())

pyrefly/lib/solver/subset.rs

@@ -79,6 +79,7 @@ use crate::types::type_var::Variance;
 use crate::types::types::Forallable;
 use crate::types::types::TArgs;
 use crate::types::types::Type;
+use crate::types::types::Var;
 
 /// Extract a `TypeAliasData` reference from a `Type` that wraps one,
 /// either directly as `Type::TypeAlias` or inside `Type::Forall`.
@@ -1523,20 +1524,25 @@ impl<'solver, 'subset, Ans: LookupAnswer> Subset<'solver, 'subset, Ans> {
         // parameters that the branches are matched against. Otherwise, the branches see fully
         // known argument types, so the first matching branch is the one that applies.
         let mut signatures = want.toplevel_callable_signatures().peekable();
-        if signatures.peek().is_some()
-            && !signatures.any(|(callable, _)| {
-                let mut mentions_captured_var = false;
+        let param_vars: SmallSet<Var> = if signatures.peek().is_some() {
+            let mut param_vars = SmallSet::new();
+            for (callable, _) in signatures {
                 callable.params.visit(&mut |ty: &Type| {
-                    mentions_captured_var |= ty
-                        .collect_maybe_placeholder_vars()
-                        .iter()
-                        .any(|v| captured_vars.contains(v));
+                    param_vars.extend(
+                        ty.collect_maybe_placeholder_vars()
+                            .into_iter()
+                            .filter(|v| captured_vars.contains(v)),
+                    );
                 });
-                mentions_captured_var
-            })
-        {
-            return OverloadCapture::NotApplicable;
-        }
+            }
+            if param_vars.is_empty() {
+                return OverloadCapture::NotApplicable;
+            }
+            param_vars
+        } else {
+            // Without a signature to inspect, any captured var may appear in the parameters.
+            captured_vars.iter().copied().collect()
+        };
         let free_quantified_vars_in_call = self.active_call_context.free_quantified_vars_in_call();
         let captures = branches
             .into_iter()
@@ -1549,6 +1555,7 @@ impl<'solver, 'subset, Ans: LookupAnswer> Subset<'solver, 'subset, Ans> {
                         me.solver.extract_overload_branch(
                             branch_index,
                             &captured_vars,
+                            &param_vars,
                             &free_quantified_vars_in_call,
                         )
                     })

pyrefly/lib/test/generic_restrictions.rs

@@ -212,6 +212,67 @@ def f() -> None:
  "#,
 );
 
+// Modeled on `pd.Series(pd.NA, dtype="Int64")` under a bare `-> pd.Series` annotation.
+testcase!(
+    test_gradual_expected_type_restriction_violation_rejects_overload,
+    r#"
+from typing import Any, Self, assert_type, overload
+
+class NA: ...
+
+class Series[S: int | str]:
+    @overload
+    def __new__(cls, data: S, dtype: int) -> Self: ...
+    @overload
+    def __new__(cls, data: S, dtype: str | None = None) -> Self: ...
+    @overload
+    def __new__(cls, data: NA, dtype: str | None = None) -> Series: ...
+    def __new__(cls, data: Any, dtype: Any = None) -> Any: ...
+
+def unhinted() -> None:
+    assert_type(Series(NA(), dtype="Int64"), Series[Any])
+
+def hinted() -> Series:
+    return Series(NA(), dtype="Int64")
+ "#,
+);
+
+// A restriction violation discards the expected type, so the call reports what the unhinted call
+// reports.
+testcase!(
+    test_gradual_expected_type_restriction_violation_without_overloads,
+    r#"
+from typing import Any
+
+class Pair[T: (int, str)]:
+    def __init__(self, left: T, right: T) -> None: ...
+
+def unhinted(left: int | str, right: int | str | bytes) -> None:
+    Pair(left, right)  # E: `int | str` is not assignable to any of constraints `int`, `str`  # E: Argument `bytes | int | str` is not assignable to parameter `right`
+
+def hinted(left: int | str, right: int | str | bytes) -> Pair[Any]:
+    return Pair(left, right)  # E: `int | str` is not assignable to any of constraints `int`, `str`  # E: Argument `bytes | int | str` is not assignable to parameter `right`
+ "#,
+);
+
+// When every overload violates the restriction, the closest overload reports the violation once.
+testcase!(
+    test_gradual_expected_type_restriction_violation_in_every_overload,
+    r#"
+from typing import Self, overload
+
+class Box[S: int]:
+    @overload
+    def __new__(cls, data: S) -> Self: ...
+    @overload
+    def __new__(cls, data: S, extra: int) -> Self: ...
+    def __new__(cls, data: S, extra: int = 0) -> Self: ...
+
+def hinted() -> Box:
+    return Box("s")  # E: `str` is not assignable to upper bound `int` of type variable `S`
+ "#,
+);
+
 testcase!(
     test_any_bound_attribute_access,
     r#"

pyrefly/lib/test/overload.rs

@@ -2769,3 +2769,37 @@ assert_type(my_divmod(LHS(), 1), tuple[LHS, LHS])
 assert_type(my_divmod(LHS(), RHS()), tuple[RHS, RHS])
     "#,
 );
+
+// `R` appears only in the return type of `SupportsF.f`, but `Box` is invariant, so matching each
+// overload pins `R` from both sides. Once `y` solves `T`, the first matching overload applies.
+testcase!(
+    test_overloaded_protocol_method_with_invariant_return_var,
+    r#"
+from typing import Protocol, assert_type, overload
+
+class Box[T]:
+    x: T
+
+class SupportsF[T, R](Protocol):
+    def f(self, other: T, /) -> Box[R]: ...
+
+def call[T, R](x: SupportsF[T, R], y: T) -> R: ...
+
+def call_box[T, R](x: SupportsF[T, R], y: T) -> Box[R]: ...
+
+class A:
+    @overload
+    def f(self, other: int) -> Box[int]: ...
+    @overload
+    def f(self, other: float) -> Box[float]: ...
+    @overload
+    def f(self, other: complex) -> Box[complex]: ...
+    def f(self, other: complex) -> Box[int] | Box[float] | Box[complex]: ...
+
+def g(a: A, i: int, fl: float):
+    assert_type(call(a, i), int)
+    assert_type(call(a, fl), float)
+    assert_type(call_box(a, i), Box[int])
+    assert_type(call_box(a, fl), Box[float])
+    "#,
+);