diff --git a/crates/red_knot_python_semantic/resources/mdtest/binary/unions.md b/crates/red_knot_python_semantic/resources/mdtest/binary/unions.md new file mode 100644 index 0000000000000..0f5fb09bf5aff --- /dev/null +++ b/crates/red_knot_python_semantic/resources/mdtest/binary/unions.md @@ -0,0 +1,51 @@ +# Binary operations on union types + +Binary operations on union types are only available if they are supported for all possible +combinations of types: + +```py +def f1(i: int, u: int | None): + # error: [unsupported-operator] "Operator `+` is unsupported between objects of type `int` and `int | None`" + reveal_type(i + u) # revealed: Unknown + # error: [unsupported-operator] "Operator `+` is unsupported between objects of type `int | None` and `int`" + reveal_type(u + i) # revealed: Unknown +``` + +`int` can be added to `int`, and `str` can be added to `str`, but expressions of type `int | str` +cannot be added, because that would require addition of `int` and `str` or vice versa: + +```py +def f2(i: int, s: str, int_or_str: int | str): + i + i + s + s + # error: [unsupported-operator] "Operator `+` is unsupported between objects of type `int | str` and `int | str`" + reveal_type(int_or_str + int_or_str) # revealed: Unknown +``` + +However, if an operation is supported for all possible combinations, the result will be a union of +the possible outcomes: + +```py +from typing import Literal + +def f3(two_or_three: Literal[2, 3], a_or_b: Literal["a", "b"]): + reveal_type(two_or_three + two_or_three) # revealed: Literal[4, 5, 6] + reveal_type(two_or_three**two_or_three) # revealed: Literal[4, 8, 9, 27] + + reveal_type(a_or_b + a_or_b) # revealed: Literal["aa", "ab", "ba", "bb"] + + reveal_type(two_or_three * a_or_b) # revealed: Literal["aa", "bb", "aaa", "bbb"] +``` + +We treat a type annotation of `float` as a union of `int` and `float`, so union handling is relevant +here: + +```py +def f4(x: float, y: float): + reveal_type(x + y) # revealed: int | float + reveal_type(x - y) # revealed: int | float + reveal_type(x * y) # revealed: int | float + reveal_type(x / y) # revealed: int | float + reveal_type(x // y) # revealed: int | float + reveal_type(x % y) # revealed: int | float +``` diff --git a/crates/red_knot_python_semantic/src/types/infer.rs b/crates/red_knot_python_semantic/src/types/infer.rs index 634a1ca62854d..0b100bfbd5d2c 100644 --- a/crates/red_knot_python_semantic/src/types/infer.rs +++ b/crates/red_knot_python_semantic/src/types/infer.rs @@ -4048,6 +4048,23 @@ impl<'db> TypeInferenceBuilder<'db> { op: ast::Operator, ) -> Option> { match (left_ty, right_ty, op) { + (Type::Union(lhs_union), rhs, _) => { + let mut union = UnionBuilder::new(self.db()); + for lhs in lhs_union.elements(self.db()) { + let result = self.infer_binary_expression_type(*lhs, rhs, op)?; + union = union.add(result); + } + Some(union.build()) + } + (lhs, Type::Union(rhs_union), _) => { + let mut union = UnionBuilder::new(self.db()); + for rhs in rhs_union.elements(self.db()) { + let result = self.infer_binary_expression_type(lhs, *rhs, op)?; + union = union.add(result); + } + Some(union.build()) + } + // Non-todo Anys take precedence over Todos (as if we fix this `Todo` in the future, // the result would then become Any or Unknown, respectively). (any @ Type::Dynamic(DynamicType::Any), _, _) @@ -4179,7 +4196,6 @@ impl<'db> TypeInferenceBuilder<'db> { | Type::SubclassOf(_) | Type::Instance(_) | Type::KnownInstance(_) - | Type::Union(_) | Type::Intersection(_) | Type::AlwaysTruthy | Type::AlwaysFalsy @@ -4196,7 +4212,6 @@ impl<'db> TypeInferenceBuilder<'db> { | Type::SubclassOf(_) | Type::Instance(_) | Type::KnownInstance(_) - | Type::Union(_) | Type::Intersection(_) | Type::AlwaysTruthy | Type::AlwaysFalsy