diff --git a/pyrefly/lib/alt/function.rs b/pyrefly/lib/alt/function.rs index 078fd1ddcb..45dc363b60 100644 --- a/pyrefly/lib/alt/function.rs +++ b/pyrefly/lib/alt/function.rs @@ -933,8 +933,12 @@ impl<'ctx, 'answer, Ans: LookupAnswer> AnswersSolver<'ctx, 'answer, Ans> { // Constructors are always validated per spec. For other methods, // skip overload variants because self/cls annotations in overloads // are a valid pattern for type narrowing. + // Constructor overloads may additionally narrow the receiver to a + // subclass of the defining class, so that e.g. `Base.__new__` can + // return a more specific type when called via a subclass. let is_constructor = stmt.name.id == dunder::INIT || is_dunder_new; let should_validate = is_constructor || !def.metadata.flags.is_overload; + let allow_subclass_receiver = is_constructor && def.metadata.flags.is_overload; if should_validate { if def.metadata.flags.is_classmethod || is_dunder_new { self.validate_cls_annotation( @@ -942,6 +946,7 @@ impl<'ctx, 'answer, Ans: LookupAnswer> AnswersSolver<'ctx, 'answer, Ans> { &stmt.name.id, &callable, def.id_range(), + allow_subclass_receiver, errors, ); } else if !def.metadata.flags.is_staticmethod { @@ -950,6 +955,7 @@ impl<'ctx, 'answer, Ans: LookupAnswer> AnswersSolver<'ctx, 'answer, Ans> { &stmt.name.id, &callable, def.id_range(), + allow_subclass_receiver, errors, ); } @@ -2770,6 +2776,7 @@ impl<'ctx, 'answer, Ans: LookupAnswer> AnswersSolver<'ctx, 'answer, Ans> { /// Validate the self annotation on methods. /// The self type must be the defining class or a superclass of it. + /// If `allow_subclass_receiver` is set, a subclass of the defining class is also accepted. /// For `__init__`, class-scoped type variables should not be used in the self annotation /// (per spec: https://typing.python.org/en/latest/spec/constructors.html#init-method). fn validate_self_annotation( @@ -2778,6 +2785,7 @@ impl<'ctx, 'answer, Ans: LookupAnswer> AnswersSolver<'ctx, 'answer, Ans> { method_name: &Name, callable: &Callable, range: TextRange, + allow_subclass_receiver: bool, errors: &ErrorCollector, ) { if let Params::List(param_list) = &callable.params @@ -2792,6 +2800,8 @@ impl<'ctx, 'answer, Ans: LookupAnswer> AnswersSolver<'ctx, 'answer, Ans> { // TODO: we could consider checking structural subtyping here. if !self.type_order().has_superclass(cls, cls_ty.class_object()) && !self.type_order().is_protocol(cls_ty.class_object()) + && !(allow_subclass_receiver + && self.type_order().has_superclass(cls_ty.class_object(), cls)) { errors .error_builder( @@ -2842,6 +2852,8 @@ impl<'ctx, 'answer, Ans: LookupAnswer> AnswersSolver<'ctx, 'answer, Ans> { if let Type::ClassType(cls_ty) = &**inner && !self.type_order().has_superclass(cls, cls_ty.class_object()) && !self.type_order().is_protocol(cls_ty.class_object()) + && !(allow_subclass_receiver + && self.type_order().has_superclass(cls_ty.class_object(), cls)) { errors .error_builder( @@ -2863,12 +2875,14 @@ impl<'ctx, 'answer, Ans: LookupAnswer> AnswersSolver<'ctx, 'answer, Ans> { /// Validate the cls annotation on `__new__` and classmethods. /// The cls type must be `type[X]` where `X` is the defining class or a superclass of it. + /// If `allow_subclass_receiver` is set, a subclass of the defining class is also accepted. fn validate_cls_annotation( &self, cls: &Class, method_name: &Name, callable: &Callable, range: TextRange, + allow_subclass_receiver: bool, errors: &ErrorCollector, ) { if let Params::List(param_list) = &callable.params @@ -2885,6 +2899,10 @@ impl<'ctx, 'answer, Ans: LookupAnswer> AnswersSolver<'ctx, 'answer, Ans> { .type_order() .has_superclass(cls, inner_cls.class_object()) && !self.type_order().is_protocol(inner_cls.class_object()) + && !(allow_subclass_receiver + && self + .type_order() + .has_superclass(inner_cls.class_object(), cls)) { errors .error_builder( diff --git a/pyrefly/lib/test/constructors.rs b/pyrefly/lib/test/constructors.rs index 194cb80ce0..0e9e235e9d 100644 --- a/pyrefly/lib/test/constructors.rs +++ b/pyrefly/lib/test/constructors.rs @@ -944,6 +944,62 @@ class C[T]: "#, ); +// https://github.com/facebook/pyrefly/issues/4701 +testcase!( + test_new_overload_subclass_receiver, + r#" +from typing import assert_type, overload + +class Path: + @overload + def __new__(cls: type[AbsPath], path: str) -> AbsPath: ... + @overload + def __new__(cls: type[RelPath], path: str) -> RelPath: ... + @overload + def __new__(cls: type[Path], path: str) -> AbsPath | RelPath: ... + def __new__(cls: type[Path], path: str) -> Path: + raise NotImplementedError + +class AbsPath(Path): ... +class RelPath(Path): ... + +assert_type(Path("file"), AbsPath | RelPath) +assert_type(AbsPath("file"), AbsPath) +assert_type(RelPath("file"), RelPath) + "#, +); + +testcase!( + test_init_overload_subclass_receiver, + r#" +from typing import overload + +class Base: + @overload + def __init__(self: Sub, value: int) -> None: ... + @overload + def __init__(self: Base, value: str) -> None: ... + def __init__(self, value: int | str): + pass + +class Sub(Base): ... + +Sub(1) +Base("x") + "#, +); + +testcase!( + test_constructor_subclass_receiver_requires_overload, + r#" +class Base: + def __new__(cls: type[Sub]) -> Sub: ... # E: `__new__` method cls type `type[Sub]` refers to class `Sub`, which is not a superclass of class `Base` + def __init__(self: Sub) -> None: ... # E: `__init__` method self type `Sub` is not a superclass of class `Base` + +class Sub(Base): ... + "#, +); + testcase!( test_generic_in_generic, r#"