Skip to content
Open
Show file tree
Hide file tree
Changes from 2 commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
46 changes: 40 additions & 6 deletions pyrefly/lib/alt/function.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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<Type> {

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This fixes only half of the problem; __init__ has the same issue:

class I[T]:
    @overload
    def __init__(self: I[int], x: object) -> None: ...
    @overload
    def __init__(self: I[str], x: str) -> None: ...

i_wrong: Callable[[str], I[int]] = I[str]   # no error
reveal_type(partial(I[str], "x")())         # I[int]
reveal_type(I[str]("x"))                    # I[str] (direct calls are fine)

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));

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

nit: This works, but it's a little hard to read due to the length and the nesting. Here are some suggestions from an AI review for tightening it up:

  1. Move the filter into a shared helper. Pull it out as something like fn filter_overloads_by_receiver(&self, func: Type, receiver: &Type) -> Type. Call it from bind_dunder_new with type[C[...]] and from bind_dunder_init on m.func.as_type() with m.obj, before the return-type transform. That fixes __init__ without duplicating the logic, and keeps bind_dunder_new as small as it was.
  2. Simplify erase_tparams. Collecting quantifieds into a SmallSet, building a map and then calling subst can be replaced by one call: ty.subst_mut_fn(&mut |_| Some(any.clone())). Don't switch to q.as_gradual_type(), though. It uses the TypeVar's default, so class C[T = int] would erase to C[int] and wrongly filter the bare class.
  3. Shorten the OverloadType match. An or-pattern does it in one line: let (OverloadType::Function(f) | OverloadType::Forall(Forall { body: f, .. })) = sig;.
  4. Flatten the control flow. Two nested matches both fall through to t.clone(). A let-else on Type::Overload followed by one match Vec1::try_from_vec(..) (or a let-chain) gets rid of the duplicate fallback and the nesting.

Sketch of the helper (not compiled):

/// Drop overloads whose first parameter cannot accept `receiver`. Type parameters on both
/// sides are replaced with `Any`, so an overload is dropped only if no specialization of it
/// could accept the receiver. If every overload would be dropped, `func` is returned unchanged.
fn filter_overloads_by_receiver(&self, func: Type, receiver: &Type) -> Type {
    let Type::Overload(overload) = &func else {
        return func;
    };
    let any = self.heap.mk_any_implicit();
    let erase = |ty: &Type| {
        let mut ty = ty.clone();
        ty.subst_mut_fn(&mut |_| Some(any.clone()));
        ty
    };
    let receiver = erase(receiver);
    let signatures = overload.signatures.iter().filter(|sig| {
        let (OverloadType::Function(f) | OverloadType::Forall(Forall { body: f, .. })) = sig;
        f.signature
            .get_first_param()
            .is_none_or(|param| self.is_subset_eq(&receiver, &erase(param)))
    });
    match Vec1::try_from_vec(signatures.cloned().collect()) {
        Ok(signatures) => self.heap.mk_overload(Overload { signatures, metadata: overload.metadata.clone() }),
        Err(_) => func,
    }
}

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.
Expand Down
76 changes: 76 additions & 0 deletions pyrefly/lib/test/constructors.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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
"#,
);
Loading