Reject an overload that violates a restriction under a gradual expected type (#5240)

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

@@ -305,9 +305,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
@@ -2463,6 +2465,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 +4201,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/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#"