[ty] Distribute binary operators over union upper bounds
diff --git a/crates/ty_python_semantic/resources/mdtest/generics/legacy/functions.md b/crates/ty_python_semantic/resources/mdtest/generics/legacy/functions.md index bdb3f55..54b7d0c 100644 --- a/crates/ty_python_semantic/resources/mdtest/generics/legacy/functions.md +++ b/crates/ty_python_semantic/resources/mdtest/generics/legacy/functions.md
@@ -1136,9 +1136,38 @@ def f(value: T): reveal_type(value + 1) # revealed: float reveal_type(value + 1.0) # revealed: float - # TODO: Adding two values of the bounded type variable should be supported. + reveal_type(value + value) # revealed: float +``` + +## Binary operators with union upper bounds + +Operator results preserve the original type variable when each member returns `Self`: + +```py +from typing_extensions import Self, TypeVar + +class A: + def __add__(self, other: object) -> Self: + return self + +class B: + def __add__(self, other: object) -> Self: + return self + +T = TypeVar("T", bound=A | B) + +def add(left: T, right: T) -> T: + return left + right +``` + +The operands can belong to different members of the bound, so `str + A` must also be supported: + +```py +U = TypeVar("U", bound=A | str) + +def mixed(left: U, right: U): # error: [unsupported-operator] - reveal_type(value + value) # revealed: Unknown + left + right ``` ## All occurrences of the same typevar have the same type
diff --git a/crates/ty_python_semantic/resources/mdtest/generics/pep695/functions.md b/crates/ty_python_semantic/resources/mdtest/generics/pep695/functions.md index 9971b38..4495326 100644 --- a/crates/ty_python_semantic/resources/mdtest/generics/pep695/functions.md +++ b/crates/ty_python_semantic/resources/mdtest/generics/pep695/functions.md
@@ -1482,9 +1482,34 @@ def f[T: float](value: T): reveal_type(value + 1) # revealed: float reveal_type(value + 1.0) # revealed: float - # TODO: Adding two values of the bounded type variable should be supported. + reveal_type(value + value) # revealed: float +``` + +## Binary operators with union upper bounds + +Operator results preserve the original type variable when each member returns `Self`: + +```py +from typing_extensions import Self + +class A: + def __add__(self, other: object) -> Self: + return self + +class B: + def __add__(self, other: object) -> Self: + return self + +def add[T: A | B](left: T, right: T) -> T: + return left + right +``` + +The operands can belong to different members of the bound, so `str + A` must also be supported: + +```py +def mixed[U: A | str](left: U, right: U): # error: [unsupported-operator] - reveal_type(value + value) # revealed: Unknown + left + right ``` ## All occurrences of the same typevar have the same type
diff --git a/crates/ty_python_semantic/src/types/call.rs b/crates/ty_python_semantic/src/types/call.rs index 63edf6e..7f674d8 100644 --- a/crates/ty_python_semantic/src/types/call.rs +++ b/crates/ty_python_semantic/src/types/call.rs
@@ -161,31 +161,39 @@ /// Memoize the return type and deprecations from binary dunder resolution, without retaining /// the full call bindings or repeating overload selection at each expression. /// Returns `None` if resolution fails; callers remain responsible for call-site diagnostics. + /// Each operand pairs its lookup type with the receiver used for binding `Self`. pub(crate) fn try_call_bin_op_result( db: &'db dyn Db, env: &ProgramEnvironment<'db>, - left_ty: Type<'db>, + left_ty: (Type<'db>, Type<'db>), op: ast::Operator, - right_ty: Type<'db>, + right_ty: (Type<'db>, Type<'db>), ) -> Option<&'db BinaryOperationResult<'db>> { #[salsa::tracked(returns(ref), cycle_initial=|_, _, _, _, _, _| None, heap_size=ruff_memory_usage::heap_size)] fn try_call_bin_op_result_impl<'db>( db: &'db dyn Db, program: Program<'db>, - left_ty: Type<'db>, + left_ty: (Type<'db>, Type<'db>), op: ast::Operator, - right_ty: Type<'db>, + right_ty: (Type<'db>, Type<'db>), ) -> Option<BinaryOperationResult<'db>> { let env = &ProgramEnvironment::from_program(program); - Type::try_call_bin_op(db, env, left_ty, op, right_ty) - .ok() - .map(|bindings| BinaryOperationResult { - return_type: bindings.return_type(db, env), - deprecated_functions: bindings - .deprecated_functions(db) - .map(|(_, function)| function) - .collect(), - }) + Type::try_call_bin_op_with_receivers( + db, + env, + left_ty, + op, + right_ty, + MemberLookupPolicy::default(), + ) + .ok() + .map(|bindings| BinaryOperationResult { + return_type: bindings.return_type(db, env), + deprecated_functions: bindings + .deprecated_functions(db) + .map(|(_, function)| function) + .collect(), + }) } try_call_bin_op_result_impl(db, env.program(db), left_ty, op, right_ty).as_ref() @@ -216,6 +224,53 @@ right_ty: Type<'db>, policy: MemberLookupPolicy, ) -> Result<Bindings<'db>, CallBinOpError> { + Self::try_call_bin_op_with_receivers( + db, + env, + (left_ty, left_ty), + op, + (right_ty, right_ty), + policy, + ) + } + + /// Resolve operators using each branch's lookup type while retaining its narrowed receiver. + fn try_call_bin_op_with_receivers( + db: &'db dyn Db, + env: &ProgramEnvironment<'db>, + (left_ty, left_receiver): (Type<'db>, Type<'db>), + op: ast::Operator, + (right_ty, right_receiver): (Type<'db>, Type<'db>), + policy: MemberLookupPolicy, + ) -> Result<Bindings<'db>, CallBinOpError> { + let call_dunder = |lookup: Type<'db>, receiver, name: &str, argument| { + let mut arguments = CallArguments::positional([argument]); + if lookup == receiver { + lookup.try_call_dunder_with_policy( + db, + env, + name, + &mut arguments, + TypeContext::default(), + policy, + ) + } else { + Self::try_call_dunder_member_impl( + db, + env, + lookup.member_lookup_with_policy_and_receiver( + db, + env, + name, + policy | MemberLookupPolicy::NO_INSTANCE_FALLBACK, + Some(receiver), + ), + &mut arguments, + TypeContext::default(), + ) + } + }; + // We either want to call lhs.__op__ or rhs.__rop__. The full decision tree from // the Python spec [1] is: // @@ -243,36 +298,17 @@ left_class.member(db, env, reflected_dunder).place, ) { - let call_on_right_instance = right_ty.try_call_dunder_with_policy( - db, - env, - reflected_dunder, - &mut CallArguments::positional([left_ty]), - TypeContext::default(), - policy, - ); + let call_on_right_instance = + call_dunder(right_ty, right_receiver, reflected_dunder, left_receiver); if reflected_priority == ReflectedMethodPriority::Definitely { return Ok(call_on_right_instance.or_else(|_| { - left_ty.try_call_dunder_with_policy( - db, - env, - op.dunder(), - &mut CallArguments::positional([right_ty]), - TypeContext::default(), - policy, - ) + call_dunder(left_ty, left_receiver, op.dunder(), right_receiver) })?); } - let call_on_left_instance = left_ty.try_call_dunder_with_policy( - db, - env, - op.dunder(), - &mut CallArguments::positional([right_ty]), - TypeContext::default(), - policy, - ); + let call_on_left_instance = + call_dunder(left_ty, left_receiver, op.dunder(), right_receiver); return match (call_on_right_instance, call_on_left_instance) { (Ok(right_bindings), Ok(left_bindings)) => { @@ -293,26 +329,18 @@ } } - let call_on_left_instance = left_ty.try_call_dunder_with_policy( - db, - env, - op.dunder(), - &mut CallArguments::positional([right_ty]), - TypeContext::default(), - policy, - ); + let call_on_left_instance = + call_dunder(left_ty, left_receiver, op.dunder(), right_receiver); call_on_left_instance.or_else(|_| { if left_ty == right_ty { Err(CallBinOpError::NotSupported) } else { - Ok(right_ty.try_call_dunder_with_policy( - db, - env, + Ok(call_dunder( + right_ty, + right_receiver, op.reflected_dunder(), - &mut CallArguments::positional([left_ty]), - TypeContext::default(), - policy, + left_receiver, )?) } })
diff --git a/crates/ty_python_semantic/src/types/infer/builder/binary_expressions.rs b/crates/ty_python_semantic/src/types/infer/builder/binary_expressions.rs index af36444..3e43b4f 100644 --- a/crates/ty_python_semantic/src/types/infer/builder/binary_expressions.rs +++ b/crates/ty_python_semantic/src/types/infer/builder/binary_expressions.rs
@@ -15,9 +15,9 @@ use crate::types::tuple::{TupleSpecBuilder, TupleType}; use crate::types::typevar::TypeVarConstraints; use crate::types::{ - DynamicType, InternedConstraintSet, KnownClass, KnownInstanceType, LiteralValueTypeKind, - MemberLookupPolicy, Type, TypeContext, TypeVarBoundOrConstraints, TypedDictType, UnionBuilder, - UnionType, UnionTypeInstance, + DynamicType, InternedConstraintSet, IntersectionType, KnownClass, KnownInstanceType, + LiteralValueTypeKind, MemberLookupPolicy, Type, TypeContext, TypeVarBoundOrConstraints, + TypedDictType, UnionBuilder, UnionType, UnionTypeInstance, }; enum BinaryExpressionOperandTypes<'db> { @@ -424,17 +424,46 @@ op: ast::Operator, right_ty: Type<'db>, ) -> Option<Type<'db>> { - let result = Type::try_call_bin_op_result( - self.db(), - self.program_environment(), - left_ty, - op, - right_ty, - )?; - state - .deprecated_functions - .extend(&result.deprecated_functions); - Some(result.return_type) + let db = self.db(); + let env = self.program_environment(); + let alternatives = |ty: Type<'db>| { + let bound = if let Type::TypeVar(typevar) = ty + && let Some(TypeVarBoundOrConstraints::UpperBound(Type::Union(bound))) = + typevar.typevar(db).bound_or_constraints(db, env) + { + Some(bound) + } else { + None + }; + bound + .map(|bound| bound.elements(db)) + .unwrap_or_default() + .iter() + .map(move |element| { + ( + *element, + IntersectionType::from_two_elements(db, env, ty, *element), + ) + }) + .chain(bound.is_none().then_some((ty, ty))) + }; + + // A union upper bound permits different members for the two operands, even when + // they share a type variable. Resolve normal and reflected methods for each pair. + // Intersecting with the original operand preserves subtypes in `Self` returns. + let left_alternatives = alternatives(left_ty); + let right_alternatives = alternatives(right_ty); + let mut results = UnionBuilder::new(db, env); + for left in left_alternatives { + for right in right_alternatives.clone() { + let result = Type::try_call_bin_op_result(db, env, left, op, right)?; + state + .deprecated_functions + .extend(&result.deprecated_functions); + results.add_in_place(result.return_type); + } + } + Some(results.build()) } /// Infer the result type and collect deprecated methods for the enclosing operation.