From 5c3c0c4705dfba0ba777df384482464e29e7d86c Mon Sep 17 00:00:00 2001 From: David Peter Date: Thu, 17 Oct 2024 11:03:37 +0200 Subject: [PATCH] [red-knot] Inference for comparison of union types (#13781) ## Summary Add type inference for comparisons involving union types. For example: ```py one_or_two = 1 if flag else 2 reveal_type(one_or_two <= 2) # revealed: Literal[True] reveal_type(one_or_two <= 1) # revealed: bool reveal_type(one_or_two <= 0) # revealed: Literal[False] ``` closes #13779 ## Test Plan See `resources/mdtest/comparison/unions.md` --- .../resources/mdtest/comparison/unions.md | 76 +++++++++++++++++++ crates/red_knot_python_semantic/src/types.rs | 5 ++ .../src/types/builder.rs | 6 ++ .../src/types/infer.rs | 23 +++++- 4 files changed, 108 insertions(+), 2 deletions(-) create mode 100644 crates/red_knot_python_semantic/resources/mdtest/comparison/unions.md diff --git a/crates/red_knot_python_semantic/resources/mdtest/comparison/unions.md b/crates/red_knot_python_semantic/resources/mdtest/comparison/unions.md new file mode 100644 index 0000000000000..e7589b7fea778 --- /dev/null +++ b/crates/red_knot_python_semantic/resources/mdtest/comparison/unions.md @@ -0,0 +1,76 @@ +# Comparison: Unions + +## Union on one side of the comparison + +Comparisons on union types need to consider all possible cases: + +```py +one_or_two = 1 if flag else 2 + +reveal_type(one_or_two <= 2) # revealed: Literal[True] +reveal_type(one_or_two <= 1) # revealed: bool +reveal_type(one_or_two <= 0) # revealed: Literal[False] + +reveal_type(2 >= one_or_two) # revealed: Literal[True] +reveal_type(1 >= one_or_two) # revealed: bool +reveal_type(0 >= one_or_two) # revealed: Literal[False] + +reveal_type(one_or_two < 1) # revealed: Literal[False] +reveal_type(one_or_two < 2) # revealed: bool +reveal_type(one_or_two < 3) # revealed: Literal[True] + +reveal_type(one_or_two > 0) # revealed: Literal[True] +reveal_type(one_or_two > 1) # revealed: bool +reveal_type(one_or_two > 2) # revealed: Literal[False] + +reveal_type(one_or_two == 3) # revealed: Literal[False] +reveal_type(one_or_two == 1) # revealed: bool + +reveal_type(one_or_two != 3) # revealed: Literal[True] +reveal_type(one_or_two != 1) # revealed: bool + +a_or_ab = "a" if flag else "ab" + +reveal_type(a_or_ab in "ab") # revealed: Literal[True] +reveal_type("a" in a_or_ab) # revealed: Literal[True] + +reveal_type("c" not in a_or_ab) # revealed: Literal[True] +reveal_type("a" not in a_or_ab) # revealed: Literal[False] + +reveal_type("b" in a_or_ab) # revealed: bool +reveal_type("b" not in a_or_ab) # revealed: bool + +one_or_none = 1 if flag else None + +reveal_type(one_or_none is None) # revealed: bool +reveal_type(one_or_none is not None) # revealed: bool +``` + +## Union on both sides of the comparison + +With unions on both sides, we need to consider the full cross product of +options when building the resulting (union) type: + +```py +small = 1 if flag_s else 2 +large = 2 if flag_l else 3 + +reveal_type(small <= large) # revealed: Literal[True] +reveal_type(small >= large) # revealed: bool + +reveal_type(small < large) # revealed: bool +reveal_type(small > large) # revealed: Literal[False] +``` + +## Unsupported operations + +Make sure we emit a diagnostic if *any* of the possible comparisons is +unsupported. For now, we fall back to `bool` for the result type instead of +trying to infer something more precise from the other (supported) variants: + +```py +x = [1, 2] if flag else 1 + +result = 1 in x # error: "Operator `in` is not supported" +reveal_type(result) # revealed: bool +``` diff --git a/crates/red_knot_python_semantic/src/types.rs b/crates/red_knot_python_semantic/src/types.rs index cf787e0462699..067c5f4cad599 100644 --- a/crates/red_knot_python_semantic/src/types.rs +++ b/crates/red_knot_python_semantic/src/types.rs @@ -415,6 +415,11 @@ impl<'db> Type<'db> { (_, Type::Unknown | Type::Any | Type::Todo) => false, (Type::Never, _) => true, (_, Type::Never) => false, + (Type::BooleanLiteral(_), Type::Instance(class)) + if class.is_known(db, KnownClass::Bool) => + { + true + } (Type::IntLiteral(_), Type::Instance(class)) if class.is_known(db, KnownClass::Int) => { true } diff --git a/crates/red_knot_python_semantic/src/types/builder.rs b/crates/red_knot_python_semantic/src/types/builder.rs index 56b6d9a4a4be6..013e6988cccbc 100644 --- a/crates/red_knot_python_semantic/src/types/builder.rs +++ b/crates/red_knot_python_semantic/src/types/builder.rs @@ -374,6 +374,12 @@ mod tests { let union = UnionType::from_elements(&db, [t0, t1, t2, t3]).expect_union(); assert_eq!(union.elements(&db), &[bool_instance_ty, t3]); + + let result_ty = UnionType::from_elements(&db, [bool_instance_ty, t0]); + assert_eq!(result_ty, bool_instance_ty); + + let result_ty = UnionType::from_elements(&db, [t0, bool_instance_ty]); + assert_eq!(result_ty, bool_instance_ty); } #[test] diff --git a/crates/red_knot_python_semantic/src/types/infer.rs b/crates/red_knot_python_semantic/src/types/infer.rs index c6a6cf490706f..e815dba0f6432 100644 --- a/crates/red_knot_python_semantic/src/types/infer.rs +++ b/crates/red_knot_python_semantic/src/types/infer.rs @@ -57,7 +57,7 @@ use crate::types::{ }; use crate::Db; -use super::KnownClass; +use super::{KnownClass, UnionBuilder}; /// Infer all types for a [`ScopeId`], including all definitions and expressions in that scope. /// Use when checking a scope, or needing to provide a type for an arbitrary expression in the @@ -2711,6 +2711,21 @@ impl<'db> TypeInferenceBuilder<'db> { // - `[ast::CompOp::Is]`: return `false` if unequal, `bool` if equal // - `[ast::CompOp::IsNot]`: return `true` if unequal, `bool` if equal match (left, right) { + (Type::Union(union), other) => { + let mut builder = UnionBuilder::new(self.db); + for element in union.elements(self.db) { + builder = builder.add(self.infer_binary_type_comparison(*element, op, other)?); + } + Some(builder.build()) + } + (other, Type::Union(union)) => { + let mut builder = UnionBuilder::new(self.db); + for element in union.elements(self.db) { + builder = builder.add(self.infer_binary_type_comparison(other, op, *element)?); + } + Some(builder.build()) + } + (Type::IntLiteral(n), Type::IntLiteral(m)) => match op { ast::CmpOp::Eq => Some(Type::BooleanLiteral(n == m)), ast::CmpOp::NotEq => Some(Type::BooleanLiteral(n != m)), @@ -2902,6 +2917,7 @@ impl<'db> TypeInferenceBuilder<'db> { } } } + // Lookup the rich comparison `__dunder__` methods on instances (Type::Instance(left_class_ty), Type::Instance(right_class_ty)) => match op { ast::CmpOp::Lt => { @@ -2911,7 +2927,10 @@ impl<'db> TypeInferenceBuilder<'db> { _ => Some(Type::Todo), }, // TODO: handle more types - _ => Some(Type::Todo), + _ => match op { + ast::CmpOp::Is | ast::CmpOp::IsNot => Some(KnownClass::Bool.to_instance(self.db)), + _ => Some(Type::Todo), + }, } }