[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.
This commit is contained in:
Ibraheem Ahmed 2025-11-14 15:25:47 -05:00 committed by GitHub
parent ffb7bdd595
commit 2a2b719f00
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
4 changed files with 93 additions and 46 deletions

View file

@ -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

View file

@ -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
};

View file

@ -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,

View file

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