[red-knot] Decompose bool to Literal[True, False] in unions and intersections

Author: AlexWaygood

crates/red_knot_python_semantic/resources/mdtest/comparison/tuples.md

@@ -58,7 +58,22 @@ reveal_type(c >= d)  # revealed: Literal[True]
 #### Results with Ambiguity
 
 ```py
-def _(x: bool, y: int):
+class P:
+    def __lt__(self, other: "P") -> bool:
+        return True
+
+    def __le__(self, other: "P") -> bool:
+        return True
+
+    def __gt__(self, other: "P") -> bool:
+        return True
+
+    def __ge__(self, other: "P") -> bool:
+        return True
+
+class Q(P): ...
+
+def _(x: P, y: Q):
     a = (x,)
     b = (y,)
 

crates/red_knot_python_semantic/resources/mdtest/exception/control_flow.md

@@ -455,9 +455,9 @@ else:
     reveal_type(x)  # revealed: slice
 finally:
     # TODO: should be `Literal[1] | str | bytes | bool | memoryview | float | range | slice`
-    reveal_type(x)  # revealed: bool | float | slice
+    reveal_type(x)  # revealed: bool | slice | float
 
-reveal_type(x)  # revealed: bool | float | slice
+reveal_type(x)  # revealed: bool | slice | float
 ```
 
 ## Nested `try`/`except` blocks
@@ -534,7 +534,7 @@ try:
         reveal_type(x)  # revealed: slice
     finally:
         # TODO: should be `Literal[1] | str | bytes | bool | memoryview | float | range | slice`
-        reveal_type(x)  # revealed: bool | float | slice
+        reveal_type(x)  # revealed: bool | slice | float
         x = 2
         reveal_type(x)  # revealed: Literal[2]
     reveal_type(x)  # revealed: Literal[2]

crates/red_knot_python_semantic/resources/mdtest/narrow/truthiness.md

@@ -21,22 +21,22 @@ else:
 if x and not x:
     reveal_type(x)  # revealed: Never
 else:
-    reveal_type(x)  # revealed: Literal[0, -1, "", "foo", b"", b"bar"] | bool | None | tuple[()]
+    reveal_type(x)  # revealed: Literal[0, -1, "", "foo", b"", b"bar"] | bool | tuple[()] | None
 
 if not (x and not x):
-    reveal_type(x)  # revealed: Literal[0, -1, "", "foo", b"", b"bar"] | bool | None | tuple[()]
+    reveal_type(x)  # revealed: Literal[0, -1, "", "foo", b"", b"bar"] | bool | tuple[()] | None
 else:
     reveal_type(x)  # revealed: Never
 
 if x or not x:
-    reveal_type(x)  # revealed: Literal[0, -1, "", "foo", b"", b"bar"] | bool | None | tuple[()]
+    reveal_type(x)  # revealed: Literal[0, -1, "", "foo", b"", b"bar"] | bool | tuple[()] | None
 else:
     reveal_type(x)  # revealed: Never
 
 if not (x or not x):
     reveal_type(x)  # revealed: Never
 else:
-    reveal_type(x)  # revealed: Literal[0, -1, "", "foo", b"", b"bar"] | bool | None | tuple[()]
+    reveal_type(x)  # revealed: Literal[0, -1, "", "foo", b"", b"bar"] | bool | tuple[()] | None
 
 if (isinstance(x, int) or isinstance(x, str)) and x:
     reveal_type(x)  # revealed: Literal[-1, True, "foo"]

crates/red_knot_python_semantic/resources/mdtest/type_properties/is_equivalent_to.md

@@ -84,4 +84,56 @@ static_assert(
 )
 ```
 
+## Unions containing tuples containing `bool`
+
+```py
+from knot_extensions import is_equivalent_to, static_assert
+from typing_extensions import Never, Literal
+
+class P: ...
+class Q: ...
+
+static_assert(is_equivalent_to(tuple[Literal[True, False]] | P, tuple[bool] | P))
+static_assert(is_equivalent_to(P | tuple[bool], P | tuple[Literal[True, False]]))
+
+static_assert(
+    is_equivalent_to(
+        tuple[tuple[tuple[P | Q]]] | P,
+        tuple[tuple[tuple[Q | P]]] | P,
+    )
+)
+```
+
+## Unions and intersections involving `AlwaysTruthy`, `bool` and `AlwaysFalsy`
+
+```py
+from knot_extensions import AlwaysTruthy, AlwaysFalsy, static_assert, is_equivalent_to, Not
+from typing_extensions import Literal
+
+static_assert(is_equivalent_to(AlwaysTruthy | bool, Literal[False] | AlwaysTruthy))
+static_assert(is_equivalent_to(AlwaysFalsy | bool, Literal[True] | AlwaysFalsy))
+static_assert(is_equivalent_to(Not[AlwaysTruthy] | bool, Not[AlwaysTruthy] | Literal[True]))
+static_assert(is_equivalent_to(Not[AlwaysFalsy] | bool, Literal[False] | Not[AlwaysFalsy]))
+```
+
+## Unions and intersections involving `AlwaysTruthy`, `LiteralString` and `AlwaysFalsy`
+
+```py
+from knot_extensions import AlwaysTruthy, AlwaysFalsy, static_assert, is_equivalent_to, Not, Intersection
+from typing_extensions import Literal, LiteralString
+
+# TODO: these should all pass!
+
+# error: [static-assert-error]
+static_assert(is_equivalent_to(AlwaysTruthy | LiteralString, Literal[""] | AlwaysTruthy))
+# error: [static-assert-error]
+static_assert(is_equivalent_to(AlwaysFalsy | LiteralString, Intersection[LiteralString, Not[Literal[""]]] | AlwaysFalsy))
+# error: [static-assert-error]
+static_assert(is_equivalent_to(Not[AlwaysFalsy] | LiteralString, Literal[""] | Not[AlwaysFalsy]))
+# error: [static-assert-error]
+static_assert(
+    is_equivalent_to(Not[AlwaysTruthy] | LiteralString, Not[AlwaysTruthy] | Intersection[LiteralString, Not[Literal[""]]])
+)
+```
+
 [the equivalence relation]: https://typing.readthedocs.io/en/latest/spec/glossary.html#term-equivalent

crates/red_knot_python_semantic/src/types.rs

@@ -811,6 +811,47 @@ impl<'db> Type<'db> {
         }
     }
 
+    #[must_use]
+    pub fn with_normalized_bools(self, db: &'db dyn Db) -> Self {
+        const LITERAL_BOOLS: [Type; 2] = [Type::BooleanLiteral(false), Type::BooleanLiteral(true)];
+
+        match self {
+            Type::Instance(InstanceType { class }) if class.is_known(db, KnownClass::Bool) => {
+                Type::Union(UnionType::new(db, Box::from(LITERAL_BOOLS)))
+            }
+            // TODO: decompose `LiteralString` into `Literal[""] | TruthyLiteralString`?
+            // We'd need to rename this method... --Alex
+            _ => self,
+        }
+    }
+
+    #[must_use]
+    pub fn with_sorted_unions(self, db: &'db dyn Db) -> Self {
+        match self {
+            Type::Union(union) => Type::Union(union.to_sorted_union(db)),
+            Type::Intersection(intersection) => {
+                Type::Intersection(intersection.to_sorted_intersection(db))
+            }
+            Type::Tuple(tuple) => Type::Tuple(tuple.with_sorted_unions(db)),
+            Type::LiteralString
+            | Type::Instance(_)
+            | Type::AlwaysFalsy
+            | Type::AlwaysTruthy
+            | Type::BooleanLiteral(_)
+            | Type::SliceLiteral(_)
+            | Type::BytesLiteral(_)
+            | Type::StringLiteral(_)
+            | Type::Dynamic(_)
+            | Type::Never
+            | Type::FunctionLiteral(_)
+            | Type::ModuleLiteral(_)
+            | Type::ClassLiteral(_)
+            | Type::KnownInstance(_)
+            | Type::IntLiteral(_)
+            | Type::SubclassOf(_) => self,
+        }
+    }
+
     /// Return true if this type is a [subtype of] type `target`.
     ///
     /// This method returns `false` if either `self` or `other` is not fully static.
@@ -830,6 +871,12 @@ impl<'db> Type<'db> {
             return true;
         }
 
+        let left = self.with_normalized_bools(db);
+        let right = target.with_normalized_bools(db);
+        if left != self || right != target {
+            return left.is_subtype_of(db, right);
+        }
+
         // Non-fully-static types do not participate in subtyping.
         //
         // Type `A` can only be a subtype of type `B` if the set of possible runtime objects
@@ -932,7 +979,7 @@ impl<'db> Type<'db> {
                 KnownClass::Str.to_instance(db).is_subtype_of(db, target)
             }
             (Type::BooleanLiteral(_), _) => {
-                KnownClass::Bool.to_instance(db).is_subtype_of(db, target)
+                KnownClass::Int.to_instance(db).is_subtype_of(db, target)
             }
             (Type::IntLiteral(_), _) => KnownClass::Int.to_instance(db).is_subtype_of(db, target),
             (Type::BytesLiteral(_), _) => {
@@ -1048,6 +1095,11 @@ impl<'db> Type<'db> {
         if self.is_gradual_equivalent_to(db, target) {
             return true;
         }
+        let normalized_self = self.with_normalized_bools(db);
+        let normalized_target = target.with_normalized_bools(db);
+        if normalized_self != self || normalized_target != target {
+            return normalized_self.is_assignable_to(db, normalized_target);
+        }
         match (self, target) {
             // Never can be assigned to any type.
             (Type::Never, _) => true,
@@ -1148,13 +1200,20 @@ impl<'db> Type<'db> {
     pub(crate) fn is_equivalent_to(self, db: &'db dyn Db, other: Type<'db>) -> bool {
         // TODO equivalent but not identical types: TypedDicts, Protocols, type aliases, etc.
 
+        let normalized_self = self.with_normalized_bools(db);
+        let normalized_other = other.with_normalized_bools(db);
+
+        if normalized_self != self || normalized_other != other {
+            return normalized_self.is_equivalent_to(db, normalized_other);
+        }
+
         match (self, other) {
             (Type::Union(left), Type::Union(right)) => left.is_equivalent_to(db, right),
             (Type::Intersection(left), Type::Intersection(right)) => {
                 left.is_equivalent_to(db, right)
             }
             (Type::Tuple(left), Type::Tuple(right)) => left.is_equivalent_to(db, right),
-            _ => self.is_fully_static(db) && other.is_fully_static(db) && self == other,
+            _ => self == other && self.is_fully_static(db) && other.is_fully_static(db),
         }
     }
 
@@ -1189,6 +1248,13 @@ impl<'db> Type<'db> {
     ///
     /// [Summary of type relations]: https://typing.readthedocs.io/en/latest/spec/concepts.html#summary-of-type-relations
     pub(crate) fn is_gradual_equivalent_to(self, db: &'db dyn Db, other: Type<'db>) -> bool {
+        let normalized_self = self.with_normalized_bools(db);
+        let normalized_other = other.with_normalized_bools(db);
+
+        if normalized_self != self || normalized_other != other {
+            return normalized_self.is_gradual_equivalent_to(db, normalized_other);
+        }
+
         if self == other {
             return true;
         }
@@ -1221,6 +1287,12 @@ impl<'db> Type<'db> {
     /// Note: This function aims to have no false positives, but might return
     /// wrong `false` answers in some cases.
     pub(crate) fn is_disjoint_from(self, db: &'db dyn Db, other: Type<'db>) -> bool {
+        let normalized_self = self.with_normalized_bools(db);
+        let normalized_other = other.with_normalized_bools(db);
+        if normalized_self != self || normalized_other != other {
+            return normalized_self.is_disjoint_from(db, normalized_other);
+        }
+
         match (self, other) {
             (Type::Never, _) | (_, Type::Never) => true,
 
@@ -4352,12 +4424,11 @@ impl<'db> UnionType<'db> {
     /// Create a new union type with the elements sorted according to a canonical ordering.
     #[must_use]
     pub fn to_sorted_union(self, db: &'db dyn Db) -> Self {
-        let mut new_elements = self.elements(db).to_vec();
-        for element in &mut new_elements {
-            if let Type::Intersection(intersection) = element {
-                intersection.sort(db);
-            }
-        }
+        let mut new_elements: Vec<Type<'db>> = self
+            .elements(db)
+            .iter()
+            .map(|element| element.with_sorted_unions(db))
+            .collect();
         new_elements.sort_unstable_by(union_elements_ordering);
         UnionType::new(db, new_elements.into_boxed_slice())
     }
@@ -4453,19 +4524,24 @@ impl<'db> IntersectionType<'db> {
     /// according to a canonical ordering.
     #[must_use]
     pub fn to_sorted_intersection(self, db: &'db dyn Db) -> Self {
-        let mut positive = self.positive(db).clone();
-        positive.sort_unstable_by(union_elements_ordering);
-
-        let mut negative = self.negative(db).clone();
-        negative.sort_unstable_by(union_elements_ordering);
+        fn normalized_set<'db>(
+            db: &'db dyn Db,
+            elements: &FxOrderSet<Type<'db>>,
+        ) -> FxOrderSet<Type<'db>> {
+            let mut elements: FxOrderSet<Type<'db>> = elements
+                .iter()
+                .map(|ty| ty.with_sorted_unions(db))
+                .collect();
 
-        IntersectionType::new(db, positive, negative)
-    }
+            elements.sort_unstable_by(union_elements_ordering);
+            elements
+        }
 
-    /// Perform an in-place sort of this [`IntersectionType`] instance
-    /// according to a canonical ordering.
-    fn sort(&mut self, db: &'db dyn Db) {
-        *self = self.to_sorted_intersection(db);
+        IntersectionType::new(
+            db,
+            normalized_set(db, self.positive(db)),
+            normalized_set(db, self.negative(db)),
+        )
     }
 
     pub fn is_fully_static(self, db: &'db dyn Db) -> bool {
@@ -4591,23 +4667,34 @@ pub struct TupleType<'db> {
 }
 
 impl<'db> TupleType<'db> {
-    pub fn from_elements<T: Into<Type<'db>>>(
-        db: &'db dyn Db,
-        types: impl IntoIterator<Item = T>,
-    ) -> Type<'db> {
+    pub fn from_elements<I, T>(db: &'db dyn Db, types: I) -> Type<'db>
+    where
+        I: IntoIterator<Item = T>,
+        T: Into<Type<'db>>,
+    {
         let mut elements = vec![];
 
         for ty in types {
-            let ty = ty.into();
+            let ty: Type<'db> = ty.into();
             if ty.is_never() {
                 return Type::Never;
             }
-            elements.push(ty);
+            elements.push(ty.with_normalized_bools(db));
         }
 
         Type::Tuple(Self::new(db, elements.into_boxed_slice()))
     }
 
+    #[must_use]
+    pub fn with_sorted_unions(self, db: &'db dyn Db) -> Self {
+        let elements: Box<[Type<'db>]> = self
+            .elements(db)
+            .iter()
+            .map(|ty| ty.with_sorted_unions(db))
+            .collect();
+        TupleType::new(db, elements)
+    }
+
     pub fn is_equivalent_to(self, db: &'db dyn Db, other: Self) -> bool {
         let self_elements = self.elements(db);
         let other_elements = other.elements(db);