[ty] Improve generic class constructor inference (#21442)

## Summary

We currently fail to account for the type context when inferring generic
classes constructed with `__new__`, or synthesized `__init__` for
dataclasses.

Author: ibraheemdev

crates/ty_python_semantic/resources/mdtest/assignment/annotations.md

@@ -324,25 +324,25 @@ class X[T]:
     def __init__(self, value: T):
         self.value = value
 
-a: X[int] = X(1)
-reveal_type(a)  # revealed: X[int]
+x1: X[int] = X(1)
+reveal_type(x1)  # revealed: X[int]
 
-b: X[int | None] = X(1)
-reveal_type(b)  # revealed: X[int | None]
+x2: X[int | None] = X(1)
+reveal_type(x2)  # revealed: X[int | None]
 
-c: X[int | None] | None = X(1)
-reveal_type(c)  # revealed: X[int | None]
+x3: X[int | None] | None = X(1)
+reveal_type(x3)  # revealed: X[int | None]
 
-def _[T](a: X[T]):
-    b: X[T | int] = X(a.value)
-    reveal_type(b)  # revealed: X[T@_ | int]
+def _[T](x1: X[T]):
+    x2: X[T | int] = X(x1.value)
+    reveal_type(x2)  # revealed: X[T@_ | int]
 
-d: X[Any] = X(1)
-reveal_type(d)  # revealed: X[Any]
+x4: X[Any] = X(1)
+reveal_type(x4)  # revealed: X[Any]
 
 def _(flag: bool):
-    a: X[int | None] = X(1) if flag else X(2)
-    reveal_type(a)  # revealed: X[int | None]
+    x5: X[int | None] = X(1) if flag else X(2)
+    reveal_type(x5)  # revealed: X[int | None]
 ```
 
 ```py
@@ -353,8 +353,7 @@ class Y[T]:
     value: T
 
 y1: Y[Any] = Y(value=1)
-# TODO: This should reveal `Y[Any]`.
-reveal_type(y1)  # revealed: Y[int]
+reveal_type(y1)  # revealed: Y[Any]
 ```
 
 ```py
@@ -363,8 +362,7 @@ class Z[T]:
         return super().__new__(cls)
 
 z1: Z[Any] = Z(1)
-# TODO: This should reveal `Z[Any]`.
-reveal_type(z1)  # revealed: Z[int]
+reveal_type(z1)  # revealed: Z[Any]
 ```
 
 ## PEP-604 annotations are supported

crates/ty_python_semantic/src/types.rs

@@ -1132,22 +1132,6 @@ impl<'db> Type<'db> {
         }
     }
 
-    /// If the type is a generic class constructor, returns the class instance type.
-    pub(crate) fn synthesized_constructor_return_ty(self, db: &'db dyn Db) -> Option<Type<'db>> {
-        // TODO: This does not correctly handle unions or intersections. It also does not handle
-        // constructors that are not represented as bound methods, e.g. `__new__`, or synthesized
-        // dataclass initializers.
-        if let Type::BoundMethod(method) = self
-            && let Type::NominalInstance(instance) = method.self_instance(db)
-            && method.function(db).name(db).as_str() == "__init__"
-        {
-            let class_ty = instance.class_literal(db).identity_specialization(db);
-            Some(Type::instance(db, class_ty))
-        } else {
-            None
-        }
-    }
-
     pub const fn is_property_instance(&self) -> bool {
         matches!(self, Type::PropertyInstance(..))
     }
@@ -6340,8 +6324,12 @@ impl<'db> Type<'db> {
         let new_call_outcome = new_method.and_then(|new_method| {
             match new_method.place.try_call_dunder_get(db, self_type) {
                 Place::Defined(new_method, _, boundness) => {
-                    let result =
-                        new_method.try_call(db, argument_types.with_self(Some(self_type)).as_ref());
+                    let argument_types = argument_types.with_self(Some(self_type));
+                    let result = new_method
+                        .bindings(db)
+                        .with_constructor_instance_type(init_ty)
+                        .match_parameters(db, &argument_types)
+                        .check_types(db, &argument_types, tcx, &[]);
 
                     if boundness == Definedness::PossiblyUndefined {
                         Some(Err(DunderNewCallError::PossiblyUnbound(result.err())))
@@ -6354,7 +6342,35 @@ impl<'db> Type<'db> {
         });
 
         let init_call_outcome = if new_call_outcome.is_none() || !init_method.is_undefined() {
-            Some(init_ty.try_call_dunder(db, "__init__", argument_types, tcx))
+            let call_result = match init_ty
+                .member_lookup_with_policy(
+                    db,
+                    "__init__".into(),
+                    MemberLookupPolicy::NO_INSTANCE_FALLBACK,
+                )
+                .place
+            {
+                Place::Undefined => Err(CallDunderError::MethodNotAvailable),
+                Place::Defined(dunder_callable, _, boundness) => {
+                    let bindings = dunder_callable
+                        .bindings(db)
+                        .with_constructor_instance_type(init_ty);
+
+                    bindings
+                        .match_parameters(db, &argument_types)
+                        .check_types(db, &argument_types, tcx, &[])
+                        .map_err(CallDunderError::from)
+                        .and_then(|bindings| {
+                            if boundness == Definedness::PossiblyUndefined {
+                                Err(CallDunderError::PossiblyUnbound(Box::new(bindings)))
+                            } else {
+                                Ok(bindings)
+                            }
+                        })
+                }
+            };
+
+            Some(call_result)
         } else {
             None
         };

crates/ty_python_semantic/src/types/call/bind.rs

@@ -53,6 +53,9 @@ pub(crate) struct Bindings<'db> {
     /// The type that is (hopefully) callable.
     callable_type: Type<'db>,
 
+    /// The type of the instance being constructed, if this signature is for a constructor.
+    constructor_instance_type: Option<Type<'db>>,
+
     /// By using `SmallVec`, we avoid an extra heap allocation for the common case of a non-union
     /// type.
     elements: SmallVec<[CallableBinding<'db>; 1]>,
@@ -77,6 +80,7 @@ impl<'db> Bindings<'db> {
             callable_type,
             elements,
             argument_forms: ArgumentForms::new(0),
+            constructor_instance_type: None,
         }
     }
 
@@ -89,6 +93,22 @@ impl<'db> Bindings<'db> {
         }
     }
 
+    pub(crate) fn with_constructor_instance_type(
+        mut self,
+        constructor_instance_type: Type<'db>,
+    ) -> Self {
+        self.constructor_instance_type = Some(constructor_instance_type);
+
+        for binding in &mut self.elements {
+            binding.constructor_instance_type = Some(constructor_instance_type);
+            for binding in &mut binding.overloads {
+                binding.constructor_instance_type = Some(constructor_instance_type);
+            }
+        }
+
+        self
+    }
+
     pub(crate) fn set_dunder_call_is_possibly_unbound(&mut self) {
         for binding in &mut self.elements {
             binding.dunder_call_is_possibly_unbound = true;
@@ -107,6 +127,7 @@ impl<'db> Bindings<'db> {
         Self {
             callable_type: self.callable_type,
             argument_forms: self.argument_forms,
+            constructor_instance_type: self.constructor_instance_type,
             elements: self.elements.into_iter().map(f).collect(),
         }
     }
@@ -240,6 +261,10 @@ impl<'db> Bindings<'db> {
         self.callable_type
     }
 
+    pub(crate) fn constructor_instance_type(&self) -> Option<Type<'db>> {
+        self.constructor_instance_type
+    }
+
     /// Returns the return type of the call. For successful calls, this is the actual return type.
     /// For calls with binding errors, this is a type that best approximates the return type. For
     /// types that are not callable, returns `Type::Unknown`.
@@ -1357,6 +1382,7 @@ impl<'db> From<CallableBinding<'db>> for Bindings<'db> {
             callable_type: from.callable_type,
             elements: smallvec_inline![from],
             argument_forms: ArgumentForms::new(0),
+            constructor_instance_type: None,
         }
     }
 }
@@ -1370,6 +1396,7 @@ impl<'db> From<Binding<'db>> for Bindings<'db> {
             signature_type,
             dunder_call_is_possibly_unbound: false,
             bound_type: None,
+            constructor_instance_type: None,
             overload_call_return_type: None,
             matching_overload_before_type_checking: None,
             overloads: smallvec_inline![from],
@@ -1378,6 +1405,7 @@ impl<'db> From<Binding<'db>> for Bindings<'db> {
             callable_type,
             elements: smallvec_inline![callable_binding],
             argument_forms: ArgumentForms::new(0),
+            constructor_instance_type: None,
         }
     }
 }
@@ -1409,6 +1437,9 @@ pub(crate) struct CallableBinding<'db> {
     /// The type of the bound `self` or `cls` parameter if this signature is for a bound method.
     pub(crate) bound_type: Option<Type<'db>>,
 
+    /// The type of the instance being constructed, if this signature is for a constructor.
+    pub(crate) constructor_instance_type: Option<Type<'db>>,
+
     /// The return type of this overloaded callable.
     ///
     /// This is [`Some`] only in the following cases:
@@ -1457,6 +1488,7 @@ impl<'db> CallableBinding<'db> {
             signature_type,
             dunder_call_is_possibly_unbound: false,
             bound_type: None,
+            constructor_instance_type: None,
             overload_call_return_type: None,
             matching_overload_before_type_checking: None,
             overloads,
@@ -1469,6 +1501,7 @@ impl<'db> CallableBinding<'db> {
             signature_type,
             dunder_call_is_possibly_unbound: false,
             bound_type: None,
+            constructor_instance_type: None,
             overload_call_return_type: None,
             matching_overload_before_type_checking: None,
             overloads: smallvec![],
@@ -2689,7 +2722,7 @@ struct ArgumentTypeChecker<'a, 'db> {
     arguments: &'a CallArguments<'a, 'db>,
     argument_matches: &'a [MatchedArgument<'db>],
     parameter_tys: &'a mut [Option<Type<'db>>],
-    callable_type: Type<'db>,
+    constructor_instance_type: Option<Type<'db>>,
     call_expression_tcx: TypeContext<'db>,
     return_ty: Type<'db>,
     errors: &'a mut Vec<BindingError<'db>>,
@@ -2706,7 +2739,7 @@ impl<'a, 'db> ArgumentTypeChecker<'a, 'db> {
         arguments: &'a CallArguments<'a, 'db>,
         argument_matches: &'a [MatchedArgument<'db>],
         parameter_tys: &'a mut [Option<Type<'db>>],
-        callable_type: Type<'db>,
+        constructor_instance_type: Option<Type<'db>>,
         call_expression_tcx: TypeContext<'db>,
         return_ty: Type<'db>,
         errors: &'a mut Vec<BindingError<'db>>,
@@ -2717,7 +2750,7 @@ impl<'a, 'db> ArgumentTypeChecker<'a, 'db> {
             arguments,
             argument_matches,
             parameter_tys,
-            callable_type,
+            constructor_instance_type,
             call_expression_tcx,
             return_ty,
             errors,
@@ -2759,8 +2792,7 @@ impl<'a, 'db> ArgumentTypeChecker<'a, 'db> {
         };
 
         let return_with_tcx = self
-            .callable_type
-            .synthesized_constructor_return_ty(self.db)
+            .constructor_instance_type
             .or(self.signature.return_ty)
             .zip(self.call_expression_tcx.annotation);
 
@@ -3109,6 +3141,9 @@ pub(crate) struct Binding<'db> {
     /// it may be a `__call__` method.
     pub(crate) signature_type: Type<'db>,
 
+    /// The type of the instance being constructed, if this signature is for a constructor.
+    pub(crate) constructor_instance_type: Option<Type<'db>>,
+
     /// Return type of the call.
     return_ty: Type<'db>,
 
@@ -3140,6 +3175,7 @@ impl<'db> Binding<'db> {
             signature,
             callable_type: signature_type,
             signature_type,
+            constructor_instance_type: None,
             return_ty: Type::unknown(),
             inferable_typevars: InferableTypeVars::None,
             specialization: None,
@@ -3204,7 +3240,7 @@ impl<'db> Binding<'db> {
             arguments,
             &self.argument_matches,
             &mut self.parameter_tys,
-            self.callable_type,
+            self.constructor_instance_type,
             call_expression_tcx,
             self.return_ty,
             &mut self.errors,

crates/ty_python_semantic/src/types/infer/builder.rs

@@ -6527,10 +6527,7 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> {
             // TODO: Checking assignability against the full declared type could help avoid
             // cases where the constraint solver is not smart enough to solve complex unions.
             // We should see revisit this after the new constraint solver is implemented.
-            if speculated_bindings
-                .callable_type()
-                .synthesized_constructor_return_ty(db)
-                .is_none()
+            if speculated_bindings.constructor_instance_type().is_none()
                 && !speculated_bindings
                     .return_type(db)
                     .is_assignable_to(db, narrowed_ty)