Allow subclass receivers on overloaded constructors

pyrefly/lib/alt/function.rs

@@ -933,15 +933,20 @@ 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(
                         cls,
                         &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(

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#"