Handle maybe-unbound __iadd__-like operators in augmented assignments (#14044)

## Summary

One of the follow-ups from augmented assignment inference, now that
`Type::Unbound` has been removed.

---------

Co-authored-by: Alex Waygood <Alex.Waygood@Gmail.com>
This commit is contained in:
Charlie Marsh 2024-11-01 13:15:35 -04:00 committed by GitHub
parent 099f077311
commit 487941ea66
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 29 additions and 9 deletions

View file

@ -1550,14 +1550,11 @@ impl<'db> TypeInferenceBuilder<'db> {
// If the target defines, e.g., `__iadd__`, infer the augmented assignment as a call to that
// dunder.
if let Type::Instance(class) = target_type {
if let Some(class_member) = class.class_member(self.db, op.in_place_dunder()).as_type()
if let Symbol::Type(class_member, boundness) =
class.class_member(self.db, op.in_place_dunder())
{
// TODO: Handle the case where boundness is `MayBeUnbound`: fall back
// to the binary-op behavior below and union the result with calling
// the possibly-unbound in-place dunder.
let call = class_member.call(self.db, &[target_type, value_type]);
return match call.return_ty_result(
let augmented_return_ty = match call.return_ty_result(
self.db,
AnyNodeRef::StmtAugAssign(assignment),
&mut self.diagnostics,
@ -1576,6 +1573,30 @@ impl<'db> TypeInferenceBuilder<'db> {
e.return_ty()
}
};
return match boundness {
Boundness::Bound => augmented_return_ty,
Boundness::MayBeUnbound => {
let left_ty = target_type;
let right_ty = value_type;
let binary_return_ty = self.infer_binary_expression_type(left_ty, right_ty, *op)
.unwrap_or_else(|| {
self.diagnostics.add(
assignment.into(),
"unsupported-operator",
format_args!(
"Operator `{op}=` is unsupported between objects of type `{}` and `{}`",
left_ty.display(self.db),
right_ty.display(self.db)
),
);
Type::Unknown
});
UnionType::from_elements(self.db, [augmented_return_ty, binary_return_ty])
}
};
}
}