diff --git a/pyrefly/lib/alt/function.rs b/pyrefly/lib/alt/function.rs index 078fd1ddcb..caf46328d6 100644 --- a/pyrefly/lib/alt/function.rs +++ b/pyrefly/lib/alt/function.rs @@ -2569,13 +2569,47 @@ impl<'ctx, 'answer, Ans: LookupAnswer> AnswersSolver<'ctx, 'answer, Ans> { self.bind_bound_method_type(&m.func, &m.obj, skip_instantiation, is_subset) } + /// Bind `__new__` to `cls`, dropping overloads whose first parameter (`cls`) cannot accept + /// the class object. Type parameters are treated as `Any` on both sides, so overloads that + /// depend on their own type parameters, or on the free type parameters of a bare generic + /// class, are kept. If every overload would be dropped, the original overload is kept so + /// that the callable stays usable. pub fn bind_dunder_new(&self, t: &Type, cls: ClassType) -> Option { - self.bind_function( - t, - &self.heap.mk_type_of(self.heap.mk_class_type(cls)), - false, - &mut |a, b| self.is_subset_eq(a, b), - ) + let cls_type = self.heap.mk_type_of(self.heap.mk_class_type(cls)); + let any = self.heap.mk_any_implicit(); + let erase_tparams = |ty: &Type| { + let mut qs = SmallSet::new(); + ty.collect_quantifieds(&mut qs); + ty.clone().subst(&qs.iter().map(|q| (*q, &any)).collect()) + }; + let t = match t { + Type::Overload(overload) => { + let receiver = erase_tparams(&cls_type); + let applicable: Vec<_> = overload + .signatures + .iter() + .filter(|sig| { + let func = match sig { + OverloadType::Function(f) => f, + OverloadType::Forall(forall) => &forall.body, + }; + func.signature.get_first_param().is_none_or(|cls_param| { + self.is_subset_eq(&receiver, &erase_tparams(cls_param)) + }) + }) + .cloned() + .collect(); + match Vec1::try_from_vec(applicable) { + Ok(signatures) => self.heap.mk_overload(Overload { + signatures, + metadata: overload.metadata.clone(), + }), + Err(_) => t.clone(), + } + } + _ => t.clone(), + }; + self.bind_function(&t, &cls_type, false, &mut |a, b| self.is_subset_eq(a, b)) } /// Bind a `__init__` method for constructor callable conversion. diff --git a/pyrefly/lib/test/constructors.rs b/pyrefly/lib/test/constructors.rs index 194cb80ce0..80df8c0c26 100644 --- a/pyrefly/lib/test/constructors.rs +++ b/pyrefly/lib/test/constructors.rs @@ -1406,3 +1406,79 @@ class Foo: Foo(untyped(1), "bad") # E: Argument `Literal['bad']` is not assignable to parameter `b` with type `int` "#, ); + +// https://github.com/facebook/pyrefly/issues/5053 +testcase!( + test_overloaded_new_cls_annotation_as_callable, + r#" +from collections.abc import Callable +from functools import partial +from typing import Any, assert_type, overload + +class C[T]: + @overload + def __new__(cls: type[C[int]], x: object) -> C[int]: ... + @overload + def __new__(cls: type[C[str]], x: str) -> C[str]: ... + def __new__(cls: type[Any], x: object) -> C[int] | C[str]: + raise NotImplementedError + +ok: Callable[[str], C[str]] = C[str] +wrong: Callable[[str], C[int]] = C[str] # E: `type[C[str]]` is not assignable to `(str) -> C[int]` +assert_type(partial(C[str], "x")(), C[str]) + "#, +); + +testcase!( + test_overloaded_new_cls_annotation_as_callable_variants, + r#" +from collections.abc import Callable +from typing import Any, assert_type, overload + +# Overloads without a `cls` annotation are all kept. +class NoAnnotation[T]: + @overload + def __new__(cls, x: int) -> "NoAnnotation[int]": ... + @overload + def __new__(cls, x: str) -> "NoAnnotation[str]": ... + def __new__(cls, x: Any) -> Any: ... + +a: Callable[[int], NoAnnotation[int]] = NoAnnotation +b: Callable[[str], NoAnnotation[str]] = NoAnnotation + +# A bare generic class keeps every overload, since its type parameter is not yet solved. +class C[T]: + @overload + def __new__(cls: type[C[int]], x: object) -> C[int]: ... + @overload + def __new__(cls: type[C[str]], x: str) -> C[str]: ... + def __new__(cls: type[Any], x: object) -> C[int] | C[str]: ... + +bare: Callable[[str], C[str]] = C + +# Both overloads match `C[Any]`. +def f(c: type[C[Any]]): + g: Callable[[str], C[str]] = c + +# A `cls` annotation that is a type variable accepts every specialization. +class Generic_[T]: + @overload + def __new__[S](cls: type[S], x: int) -> S: ... + @overload + def __new__[S](cls: type[S], x: str) -> S: ... + def __new__[S](cls: type[S], x: Any) -> S: ... + +h: Callable[[int], Generic_[str]] = Generic_[str] +i: Callable[[str], Generic_[str]] = Generic_[str] + +# `__init__`-only classes are unaffected. +class InitOnly[T]: + def __init__(self, x: T) -> None: ... + +j: Callable[[int], InitOnly[int]] = InitOnly[int] + +# Inherited overloaded `__new__` is filtered against the subclass. +class D(C[str]): ... +k: Callable[[str], C[str]] = D + "#, +);