Convert type[Self] to callable from class constructor

pyrefly/lib/alt/call.rs

@@ -3020,6 +3020,24 @@ impl<'ctx, 'answer, Ans: LookupAnswer> AnswersSolver<'ctx, 'answer, Ans> {
         self.normalize_class_constructor_tparams(constructor, class_tparams)
     }
 
+    /// Convert `type[Self]` to a callable. As in a `type[Self]` call, a constructor that returns
+    /// an instance of the class returns `Self` instead. A type variable bound to the class is
+    /// converted the same way, with the type variable as `self_type`.
+    pub fn constructor_to_callable_for_self_type(&self, cls: &ClassType, self_type: &Type) -> Type {
+        let return_self = |mut ty: Type| {
+            ty.transform_toplevel_callable_signatures(|callable: &mut Callable, _| {
+                if self.is_compatible_constructor_return(&callable.ret, cls.class_object()) {
+                    callable.ret = self_type.clone();
+                }
+            });
+            ty
+        };
+        match self.constructor_to_callable(cls) {
+            Type::Union(union) => self.unions(union.members.into_iter().map(return_self).collect()),
+            ty => return_self(ty),
+        }
+    }
+
     /// Normalize class type parameters for each callable branch in `ty`.
     /// Normalization sets type parameters that don't appear in the callable to their gradual
     /// fallback and makes the callable generic over the type parameters that do appear.

pyrefly/lib/alt/function.rs

@@ -2867,7 +2867,10 @@ impl<'ctx, 'answer, Ans: LookupAnswer> AnswersSolver<'ctx, 'answer, Ans> {
     pub fn bind_dunder_new(&self, t: &Type, cls: ClassType) -> Option<Type> {
         let cls_type = self.heap.mk_type_of(self.heap.mk_class_type(cls));
         let t = self.filter_overloads_by_receiver(t.clone(), &cls_type);
-        self.bind_function(&t, &cls_type, false, &mut |a, b| self.is_subset_eq(a, b))
+        let mut bound =
+            self.bind_function(&t, &cls_type, false, &mut |a, b| self.is_subset_eq(a, b))?;
+        self.expand_mut(&mut bound);
+        Some(bound)
     }
 
     /// Bind a `__init__` method for constructor callable conversion.

pyrefly/lib/solver/subset.rs

@@ -2597,6 +2597,21 @@ impl<'solver, 'subset, Ans: LookupAnswer> Subset<'solver, 'subset, Ans> {
             {
                 self.is_subset_eq(&self.type_order.constructor_to_callable(got_cls), want)
             }
+            (Type::Type(inner), Type::BoundMethod(_) | Type::Callable(_) | Type::Function(_))
+                if let Some(got_cls) = match &**inner {
+                    Type::SelfType(cls) => Some(cls),
+                    Type::Quantified(q) => match q.restriction() {
+                        Restriction::Bound(Type::ClassType(cls)) => Some(cls),
+                        _ => None,
+                    },
+                    _ => None,
+                } =>
+            {
+                let constructor = self
+                    .type_order
+                    .constructor_to_callable_for_self_type(got_cls, inner);
+                self.is_subset_eq(&constructor, want)
+            }
             (Type::ClassDef(got), Type::BoundMethod(_) | Type::Callable(_) | Type::Function(_)) => {
                 let constructor = self.type_order.constructor_to_callable_for_class_def(got);
                 self.is_subset_eq(&constructor, want)

pyrefly/lib/solver/type_order.rs

@@ -219,6 +219,10 @@ impl<'solver, Ans: LookupAnswer> TypeOrder<'solver, Ans> {
         self.0.constructor_to_callable_for_class_def(cls)
     }
 
+    pub fn constructor_to_callable_for_self_type(self, cls: &ClassType, self_type: &Type) -> Type {
+        self.0.constructor_to_callable_for_self_type(cls, self_type)
+    }
+
     pub fn instantiate_fresh_forall(self, forall: Forall<Forallable>) -> (QuantifiedHandle, Type) {
         self.0.instantiate_fresh_forall(forall)
     }

pyrefly/lib/test/constructors.rs

@@ -1544,6 +1544,72 @@ ok: Callable[[int], list[C[int]]] = C
     "#,
 );
 
+// https://github.com/meta/pyrefly/issues/5591
+testcase!(
+    test_type_of_self_as_callable,
+    r#"
+from typing import Callable, reveal_type
+
+def identity[**P, R](cb: Callable[P, R]) -> Callable[P, R]:
+    return cb
+
+class C:
+    def __init__(self, x: int) -> None: ...
+
+    @classmethod
+    def method(cls) -> None:
+        reveal_type(identity(cls))  # E: revealed type: (x: int) -> Self@C
+        ok: Callable[[int], C] = cls
+        bad: Callable[[str], C] = cls  # E: `type[Self@C]` is not assignable to `(str) -> C`
+
+    @classmethod
+    def make[S: C](cls: type[S]) -> S:
+        reveal_type(identity(cls))  # E: revealed type: (x: int) -> S
+        return identity(cls)(1)
+
+class G[T]:
+    def __init__(self, x: T) -> None: ...
+
+    @classmethod
+    def method(cls) -> None:
+        reveal_type(identity(cls))  # E: revealed type: (x: T) -> Self@G
+    "#,
+);
+
+testcase!(
+    test_type_of_self_as_callable_with_type_var_new,
+    r#"
+from typing import Callable, Self, TypeVar, assert_type, reveal_type
+
+def identity[**P, R](cb: Callable[P, R]) -> Callable[P, R]:
+    return cb
+
+class C:
+    def __new__[S: C](cls: type[S], x: int) -> S:
+        raise NotImplementedError
+
+    @classmethod
+    def make(cls) -> Self:
+        reveal_type(identity(cls))  # E: revealed type: (x: int) -> Self@C
+        return identity(cls)(1)
+
+_S = TypeVar("_S", bound="Legacy")
+
+class Legacy:
+    def __new__(cls: type[_S], x: int) -> _S:
+        raise NotImplementedError
+
+    @classmethod
+    def make(cls) -> Self:
+        reveal_type(identity(cls))  # E: revealed type: (x: int) -> Self@Legacy
+        return identity(cls)(1)
+
+class D(C): ...
+
+assert_type(D.make(), D)
+    "#,
+);
+
 // Soft errors, such as the ignored-by-default `unknown-argument-type`, must not make the
 // `__init__` check look redundant.
 testcase!(