[red-knot] More precise inference for chained boolean expressions

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

@@ -34,7 +34,7 @@ x = A() < B() < C()
 reveal_type(x)  # revealed: A | B
 
 y = 0 < 1 < A() < 3
-reveal_type(y)  # revealed: bool | A
+reveal_type(y)  # revealed: Literal[False] | A
 
 z = 10 < 0 < A() < B() < C()
 reveal_type(z)  # revealed: Literal[False]

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

@@ -219,3 +219,49 @@ else:
 # TODO: It should be A. We should improve UnionBuilder or IntersectionBuilder. (issue #15023)
 reveal_type(y)  # revealed: A & ~AlwaysTruthy | A & ~AlwaysFalsy
 ```
+
+## Filtering in chained boolean expressions
+
+```py
+from typing import Literal
+
+class A: ...
+
+def _(x: Literal[0, 1]):
+    reveal_type(x or A())  # revealed: Literal[1] | A
+    reveal_type(x and A())  # revealed: Literal[0] | A
+
+# Filtering should not happen in cases where it makes the result type more complex
+def _(x: str):
+    reveal_type(x or A())  # revealed: str | A
+    reveal_type(x and A())  # revealed: str | A
+
+def _(x: bool | str):
+    reveal_type(x or A())  # revealed: Literal[True] | str | A
+    reveal_type(x and A())  # revealed: Literal[False] | str | A
+
+class Falsy:
+    def __bool__(self) -> Literal[False]: ...
+
+class Truthy:
+    def __bool__(self) -> Literal[True]: ...
+
+def _(x: Falsy | Truthy):
+    reveal_type(x or A())  # revealed: Truthy | A
+    reveal_type(x and A())  # revealed: Falsy | A
+
+class MetaFalsy(type):
+    def __bool__(self) -> Literal[False]: ...
+
+class MetaTruthy(type):
+    def __bool__(self) -> Literal[False]: ...
+
+class FalsyClass(metaclass=MetaFalsy): ...
+class TruthyClass(metaclass=MetaTruthy): ...
+
+def _(x: type[FalsyClass] | type[TruthyClass]):
+    # TODO: Should be `type[TruthyClass] | A`
+    reveal_type(x or A())  # revealed: type[FalsyClass] | type[TruthyClass] | A
+    # TODO: Should be `type[FalsyClass] | A`
+    reveal_type(x and A())  # revealed: type[FalsyClass] | type[TruthyClass] | A
+```

crates/red_knot_python_semantic/src/types/infer.rs

@@ -3573,28 +3573,61 @@ impl<'db> TypeInferenceBuilder<'db> {
         values: impl IntoIterator<Item = Type<'db>>,
         n_values: usize,
     ) -> Type<'db> {
-        let mut done = false;
-        UnionType::from_elements(
-            db,
-            values.into_iter().enumerate().map(|(i, ty)| {
-                if done {
-                    Type::Never
-                } else {
-                    let is_last = i == n_values - 1;
-                    match (ty.bool(db), is_last, op) {
-                        (Truthiness::Ambiguous, _, _) => ty,
-                        (Truthiness::AlwaysTrue, false, ast::BoolOp::And) => Type::Never,
-                        (Truthiness::AlwaysFalse, false, ast::BoolOp::Or) => Type::Never,
-                        (Truthiness::AlwaysFalse, _, ast::BoolOp::And)
-                        | (Truthiness::AlwaysTrue, _, ast::BoolOp::Or) => {
-                            done = true;
-                            ty
-                        }
-                        (_, true, _) => ty,
+        fn filter<'db>(db: &'db dyn Db, ty: Type<'db>, filtered_ty: Type<'db>) -> Type<'db> {
+            let filter_single = |ty| {
+                IntersectionBuilder::new(db)
+                    .add_positive(ty)
+                    .add_negative(filtered_ty)
+                    .build()
+            };
+
+            match ty {
+                Type::BooleanLiteral(_) => filter_single(ty),
+                Type::Instance(instance) if instance.class.is_known(db, KnownClass::Bool) => {
+                    filter_single(ty)
+                }
+                Type::Union(union) => union.map(db, |element_ty| {
+                    if element_ty.is_subtype_of(db, filtered_ty) {
+                        Type::Never
+                    } else {
+                        filter(db, *element_ty, filtered_ty)
                     }
+                }),
+                _ => ty,
+            }
+        }
+
+        let mut done = false;
+
+        let filtered_ty = match op {
+            ast::BoolOp::And => Type::AlwaysTruthy,
+            ast::BoolOp::Or => Type::AlwaysFalsy,
+        };
+
+        let elements = values.into_iter().enumerate().map(|(i, ty)| {
+            if done {
+                return Type::Never;
+            }
+
+            let is_last = i == n_values - 1;
+
+            match (ty.bool(db), is_last, op) {
+                (Truthiness::AlwaysTrue, false, ast::BoolOp::And) => Type::Never,
+                (Truthiness::AlwaysFalse, false, ast::BoolOp::Or) => Type::Never,
+
+                (Truthiness::AlwaysFalse, _, ast::BoolOp::And)
+                | (Truthiness::AlwaysTrue, _, ast::BoolOp::Or) => {
+                    done = true;
+                    ty
                 }
-            }),
-        )
+
+                (Truthiness::Ambiguous, false, _) => filter(db, ty, filtered_ty),
+
+                (_, true, _) => ty,
+            }
+        });
+
+        UnionType::from_elements(db, elements)
     }
 
     fn infer_compare_expression(&mut self, compare: &ast::ExprCompare) -> Type<'db> {