diff --git a/crates/red_knot_python_semantic/resources/mdtest/assignment/augmented.md b/crates/red_knot_python_semantic/resources/mdtest/assignment/augmented.md index e0e680b805..061c965080 100644 --- a/crates/red_knot_python_semantic/resources/mdtest/assignment/augmented.md +++ b/crates/red_knot_python_semantic/resources/mdtest/assignment/augmented.md @@ -85,7 +85,7 @@ f = Foo() # that `Foo.__iadd__` may be unbound as additional context. f += "Hello, world!" -reveal_type(f) # revealed: int +reveal_type(f) # revealed: int | @Todo ``` ## Partially bound with `__add__` @@ -104,8 +104,7 @@ class Foo: f = Foo() f += "Hello, world!" -# TODO(charlie): This should be `int | str`, since `__iadd__` may be unbound. -reveal_type(f) # revealed: int +reveal_type(f) # revealed: int | str ``` ## Partially bound target union diff --git a/crates/red_knot_python_semantic/src/types/infer.rs b/crates/red_knot_python_semantic/src/types/infer.rs index 567e615a0b..fd3f807122 100644 --- a/crates/red_knot_python_semantic/src/types/infer.rs +++ b/crates/red_knot_python_semantic/src/types/infer.rs @@ -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]) + } + }; } }