[red-knot] Narrowing For Truthiness Checks (if x or if not x) (#14687)

## Summary

Fixes #14550.

Add `AlwaysTruthy` and `AlwaysFalsy` types, representing the set of objects whose `__bool__` method can only ever return `True` or `False`, respectively, and narrow `if x` and `if not x` accordingly.


## Test Plan

- New Markdown test for truthiness narrowing `narrow/truthiness.md`
- unit tests in `types.rs` and `builders.rs` (`cargo test --package
red_knot_python_semantic --lib -- types`)

crates/red_knot_python_semantic/src/types.rs

@@ -431,6 +431,11 @@ pub enum Type<'db> {
     Union(UnionType<'db>),
     /// The set of objects in all of the types in the intersection
     Intersection(IntersectionType<'db>),
+    /// Represents objects whose `__bool__` method is deterministic:
+    /// - `AlwaysTruthy`: `__bool__` always returns `True`
+    /// - `AlwaysFalsy`: `__bool__` always returns `False`
+    AlwaysTruthy,
+    AlwaysFalsy,
     /// An integer literal
     IntLiteral(i64),
     /// A boolean literal, either `True` or `False`.
@@ -717,6 +722,15 @@ impl<'db> Type<'db> {
                         .all(|&neg_ty| self.is_disjoint_from(db, neg_ty))
             }
 
+            // Note that the definition of `Type::AlwaysFalsy` depends on the return value of `__bool__`.
+            // If `__bool__` always returns True or False, it can be treated as a subtype of `AlwaysTruthy` or `AlwaysFalsy`, respectively.
+            (left, Type::AlwaysFalsy) => matches!(left.bool(db), Truthiness::AlwaysFalse),
+            (left, Type::AlwaysTruthy) => matches!(left.bool(db), Truthiness::AlwaysTrue),
+            // Currently, the only supertype of `AlwaysFalsy` and `AlwaysTruthy` is the universal set (object instance).
+            (Type::AlwaysFalsy | Type::AlwaysTruthy, _) => {
+                target.is_equivalent_to(db, KnownClass::Object.to_instance(db))
+            }
+
             // All `StringLiteral` types are a subtype of `LiteralString`.
             (Type::StringLiteral(_), Type::LiteralString) => true,
 
@@ -1105,6 +1119,16 @@ impl<'db> Type<'db> {
                 false
             }
 
+            (Type::AlwaysTruthy, ty) | (ty, Type::AlwaysTruthy) => {
+                // `Truthiness::Ambiguous` may include `AlwaysTrue` as a subset, so it's not guaranteed to be disjoint.
+                // Thus, they are only disjoint if `ty.bool() == AlwaysFalse`.
+                matches!(ty.bool(db), Truthiness::AlwaysFalse)
+            }
+            (Type::AlwaysFalsy, ty) | (ty, Type::AlwaysFalsy) => {
+                // Similarly, they are only disjoint if `ty.bool() == AlwaysTrue`.
+                matches!(ty.bool(db), Truthiness::AlwaysTrue)
+            }
+
             (Type::KnownInstance(left), right) => {
                 left.instance_fallback(db).is_disjoint_from(db, right)
             }
@@ -1238,7 +1262,9 @@ impl<'db> Type<'db> {
             | Type::LiteralString
             | Type::BytesLiteral(_)
             | Type::SliceLiteral(_)
-            | Type::KnownInstance(_) => true,
+            | Type::KnownInstance(_)
+            | Type::AlwaysFalsy
+            | Type::AlwaysTruthy => true,
             Type::SubclassOf(SubclassOfType { base }) => matches!(base, ClassBase::Class(_)),
             Type::ClassLiteral(_) | Type::Instance(_) => {
                 // TODO: Ideally, we would iterate over the MRO of the class, check if all
@@ -1340,6 +1366,7 @@ impl<'db> Type<'db> {
                 //
                 false
             }
+            Type::AlwaysTruthy | Type::AlwaysFalsy => false,
         }
     }
 
@@ -1410,7 +1437,9 @@ impl<'db> Type<'db> {
             | Type::Todo(_)
             | Type::Union(..)
             | Type::Intersection(..)
-            | Type::LiteralString => false,
+            | Type::LiteralString
+            | Type::AlwaysTruthy
+            | Type::AlwaysFalsy => false,
         }
     }
 
@@ -1578,6 +1607,10 @@ impl<'db> Type<'db> {
                 // TODO: implement tuple methods
                 todo_type!().into()
             }
+            Type::AlwaysTruthy | Type::AlwaysFalsy => {
+                // TODO return `Callable[[], Literal[True/False]]` for `__bool__` access
+                KnownClass::Object.to_instance(db).member(db, name)
+            }
             &todo @ Type::Todo(_) => todo.into(),
         }
     }
@@ -1600,6 +1633,8 @@ impl<'db> Type<'db> {
                 // TODO: see above
                 Truthiness::Ambiguous
             }
+            Type::AlwaysTruthy => Truthiness::AlwaysTrue,
+            Type::AlwaysFalsy => Truthiness::AlwaysFalse,
             instance_ty @ Type::Instance(InstanceType { class }) => {
                 if class.is_known(db, KnownClass::NoneType) {
                     Truthiness::AlwaysFalse
@@ -1912,7 +1947,9 @@ impl<'db> Type<'db> {
             | Type::StringLiteral(_)
             | Type::SliceLiteral(_)
             | Type::Tuple(_)
-            | Type::LiteralString => Type::Unknown,
+            | Type::LiteralString
+            | Type::AlwaysTruthy
+            | Type::AlwaysFalsy => Type::Unknown,
         }
     }
 
@@ -2074,6 +2111,7 @@ impl<'db> Type<'db> {
                 ClassBase::try_from_ty(db, todo_type!("Intersection meta-type"))
                     .expect("Type::Todo should be a valid ClassBase"),
             ),
+            Type::AlwaysTruthy | Type::AlwaysFalsy => KnownClass::Type.to_instance(db),
             Type::Todo(todo) => Type::subclass_of_base(ClassBase::Todo(*todo)),
         }
     }
@@ -3558,6 +3596,8 @@ pub(crate) mod tests {
         SubclassOfAbcClass(&'static str),
         StdlibModule(CoreStdlibModule),
         SliceLiteral(i32, i32, i32),
+        AlwaysTruthy,
+        AlwaysFalsy,
     }
 
     impl Ty {
@@ -3625,6 +3665,8 @@ pub(crate) mod tests {
                     Some(stop),
                     Some(step),
                 )),
+                Ty::AlwaysTruthy => Type::AlwaysTruthy,
+                Ty::AlwaysFalsy => Type::AlwaysFalsy,
             }
         }
     }
@@ -3763,6 +3805,12 @@ pub(crate) mod tests {
     )]
     #[test_case(Ty::SliceLiteral(1, 2, 3), Ty::BuiltinInstance("slice"))]
     #[test_case(Ty::SubclassOfBuiltinClass("str"), Ty::Intersection{pos: vec![], neg: vec![Ty::None]})]
+    #[test_case(Ty::IntLiteral(1), Ty::AlwaysTruthy)]
+    #[test_case(Ty::IntLiteral(0), Ty::AlwaysFalsy)]
+    #[test_case(Ty::AlwaysTruthy, Ty::BuiltinInstance("object"))]
+    #[test_case(Ty::AlwaysFalsy, Ty::BuiltinInstance("object"))]
+    #[test_case(Ty::Never, Ty::AlwaysTruthy)]
+    #[test_case(Ty::Never, Ty::AlwaysFalsy)]
     fn is_subtype_of(from: Ty, to: Ty) {
         let db = setup_db();
         assert!(from.into_type(&db).is_subtype_of(&db, to.into_type(&db)));
@@ -3797,6 +3845,10 @@ pub(crate) mod tests {
     #[test_case(Ty::BuiltinClassLiteral("str"), Ty::SubclassOfAny)]
     #[test_case(Ty::AbcInstance("ABCMeta"), Ty::SubclassOfBuiltinClass("type"))]
     #[test_case(Ty::SubclassOfBuiltinClass("str"), Ty::BuiltinClassLiteral("str"))]
+    #[test_case(Ty::IntLiteral(1), Ty::AlwaysFalsy)]
+    #[test_case(Ty::IntLiteral(0), Ty::AlwaysTruthy)]
+    #[test_case(Ty::BuiltinInstance("str"), Ty::AlwaysTruthy)]
+    #[test_case(Ty::BuiltinInstance("str"), Ty::AlwaysFalsy)]
     fn is_not_subtype_of(from: Ty, to: Ty) {
         let db = setup_db();
         assert!(!from.into_type(&db).is_subtype_of(&db, to.into_type(&db)));
@@ -3931,6 +3983,7 @@ pub(crate) mod tests {
     #[test_case(Ty::Tuple(vec![]), Ty::BuiltinClassLiteral("object"))]
     #[test_case(Ty::SubclassOfBuiltinClass("object"), Ty::None)]
     #[test_case(Ty::SubclassOfBuiltinClass("str"), Ty::LiteralString)]
+    #[test_case(Ty::AlwaysFalsy, Ty::AlwaysTruthy)]
     fn is_disjoint_from(a: Ty, b: Ty) {
         let db = setup_db();
         let a = a.into_type(&db);
@@ -3961,6 +4014,8 @@ pub(crate) mod tests {
     #[test_case(Ty::BuiltinClassLiteral("str"), Ty::BuiltinInstance("type"))]
     #[test_case(Ty::BuiltinClassLiteral("str"), Ty::SubclassOfAny)]
     #[test_case(Ty::AbcClassLiteral("ABC"), Ty::AbcInstance("ABCMeta"))]
+    #[test_case(Ty::BuiltinInstance("str"), Ty::AlwaysTruthy)]
+    #[test_case(Ty::BuiltinInstance("str"), Ty::AlwaysFalsy)]
     fn is_not_disjoint_from(a: Ty, b: Ty) {
         let db = setup_db();
         let a = a.into_type(&db);

crates/red_knot_python_semantic/src/types/builder.rs

@@ -30,6 +30,8 @@ use crate::types::{InstanceType, IntersectionType, KnownClass, Type, UnionType};
 use crate::{Db, FxOrderSet};
 use smallvec::SmallVec;
 
+use super::Truthiness;
+
 pub(crate) struct UnionBuilder<'db> {
     elements: Vec<Type<'db>>,
     db: &'db dyn Db,
@@ -243,15 +245,22 @@ impl<'db> InnerIntersectionBuilder<'db> {
             }
         } else {
             // ~Literal[True] & bool = Literal[False]
+            // ~AlwaysTruthy & bool = Literal[False]
             if let Type::Instance(InstanceType { class }) = new_positive {
                 if class.is_known(db, KnownClass::Bool) {
-                    if let Some(&Type::BooleanLiteral(value)) = self
+                    if let Some(new_type) = self
                         .negative
                         .iter()
-                        .find(|element| element.is_boolean_literal())
+                        .find(|element| {
+                            element.is_boolean_literal()
+                                | matches!(element, Type::AlwaysFalsy | Type::AlwaysTruthy)
+                        })
+                        .map(|element| {
+                            Type::BooleanLiteral(element.bool(db) != Truthiness::AlwaysTrue)
+                        })
                     {
                         *self = Self::default();
-                        self.positive.insert(Type::BooleanLiteral(!value));
+                        self.positive.insert(new_type);
                         return;
                     }
                 }
@@ -318,15 +327,15 @@ impl<'db> InnerIntersectionBuilder<'db> {
                 // simplify the representation.
                 self.add_positive(db, ty);
             }
-            // ~Literal[True] & bool = Literal[False]
-            Type::BooleanLiteral(bool)
-                if self
-                    .positive
-                    .iter()
-                    .any(|pos| *pos == KnownClass::Bool.to_instance(db)) =>
+            // bool & ~Literal[True] = Literal[False]
+            // bool & ~AlwaysTruthy = Literal[False]
+            Type::BooleanLiteral(_) | Type::AlwaysFalsy | Type::AlwaysTruthy
+                if self.positive.contains(&KnownClass::Bool.to_instance(db)) =>
             {
                 *self = Self::default();
-                self.positive.insert(Type::BooleanLiteral(!bool));
+                self.positive.insert(Type::BooleanLiteral(
+                    new_negative.bool(db) != Truthiness::AlwaysTrue,
+                ));
             }
             _ => {
                 let mut to_remove = SmallVec::<[usize; 1]>::new();
@@ -380,7 +389,7 @@ mod tests {
     use super::{IntersectionBuilder, IntersectionType, Type, UnionType};
 
     use crate::db::tests::{setup_db, TestDb};
-    use crate::types::{global_symbol, todo_type, KnownClass, UnionBuilder};
+    use crate::types::{global_symbol, todo_type, KnownClass, Truthiness, UnionBuilder};
 
     use ruff_db::files::system_path_to_file;
     use ruff_db::system::DbWithTestSystem;
@@ -997,42 +1006,43 @@ mod tests {
         assert_eq!(ty, expected);
     }
 
-    #[test_case(true)]
-    #[test_case(false)]
-    fn build_intersection_simplify_split_bool(bool_value: bool) {
+    #[test_case(Type::BooleanLiteral(true))]
+    #[test_case(Type::BooleanLiteral(false))]
+    #[test_case(Type::AlwaysTruthy)]
+    #[test_case(Type::AlwaysFalsy)]
+    fn build_intersection_simplify_split_bool(t_splitter: Type) {
         let db = setup_db();
-
-        let t_bool = KnownClass::Bool.to_instance(&db);
-        let t_boolean_literal = Type::BooleanLiteral(bool_value);
+        let bool_value = t_splitter.bool(&db) == Truthiness::AlwaysTrue;
 
         // We add t_object in various orders (in first or second position) in
         // the tests below to ensure that the boolean simplification eliminates
         // everything from the intersection, not just `bool`.
         let t_object = KnownClass::Object.to_instance(&db);
+        let t_bool = KnownClass::Bool.to_instance(&db);
 
         let ty = IntersectionBuilder::new(&db)
             .add_positive(t_object)
             .add_positive(t_bool)
-            .add_negative(t_boolean_literal)
+            .add_negative(t_splitter)
             .build();
         assert_eq!(ty, Type::BooleanLiteral(!bool_value));
 
         let ty = IntersectionBuilder::new(&db)
             .add_positive(t_bool)
             .add_positive(t_object)
-            .add_negative(t_boolean_literal)
+            .add_negative(t_splitter)
             .build();
         assert_eq!(ty, Type::BooleanLiteral(!bool_value));
 
         let ty = IntersectionBuilder::new(&db)
             .add_positive(t_object)
-            .add_negative(t_boolean_literal)
+            .add_negative(t_splitter)
             .add_positive(t_bool)
             .build();
         assert_eq!(ty, Type::BooleanLiteral(!bool_value));
 
         let ty = IntersectionBuilder::new(&db)
-            .add_negative(t_boolean_literal)
+            .add_negative(t_splitter)
             .add_positive(t_object)
             .add_positive(t_bool)
             .build();

crates/red_knot_python_semantic/src/types/class_base.rs

@@ -70,7 +70,9 @@ impl<'db> ClassBase<'db> {
             | Type::Tuple(_)
             | Type::SliceLiteral(_)
             | Type::ModuleLiteral(_)
-            | Type::SubclassOf(_) => None,
+            | Type::SubclassOf(_)
+            | Type::AlwaysFalsy
+            | Type::AlwaysTruthy => None,
             Type::KnownInstance(known_instance) => match known_instance {
                 KnownInstanceType::TypeVar(_)
                 | KnownInstanceType::TypeAliasType(_)

crates/red_knot_python_semantic/src/types/display.rs

@@ -140,6 +140,8 @@ impl Display for DisplayRepresentation<'_> {
                 }
                 f.write_str("]")
             }
+            Type::AlwaysTruthy => f.write_str("AlwaysTruthy"),
+            Type::AlwaysFalsy => f.write_str("AlwaysFalsy"),
         }
     }
 }

crates/red_knot_python_semantic/src/types/narrow.rs

@@ -196,6 +196,7 @@ impl<'db> NarrowingConstraintsBuilder<'db> {
         is_positive: bool,
     ) -> Option<NarrowingConstraints<'db>> {
         match expression_node {
+            ast::Expr::Name(name) => Some(self.evaluate_expr_name(name, is_positive)),
             ast::Expr::Compare(expr_compare) => {
                 self.evaluate_expr_compare(expr_compare, expression, is_positive)
             }
@@ -254,6 +255,31 @@ impl<'db> NarrowingConstraintsBuilder<'db> {
         }
     }
 
+    fn evaluate_expr_name(
+        &mut self,
+        expr_name: &ast::ExprName,
+        is_positive: bool,
+    ) -> NarrowingConstraints<'db> {
+        let ast::ExprName { id, .. } = expr_name;
+
+        let symbol = self
+            .symbols()
+            .symbol_id_by_name(id)
+            .expect("Should always have a symbol for every Name node");
+        let mut constraints = NarrowingConstraints::default();
+
+        constraints.insert(
+            symbol,
+            if is_positive {
+                Type::AlwaysFalsy.negate(self.db)
+            } else {
+                Type::AlwaysTruthy.negate(self.db)
+            },
+        );
+
+        constraints
+    }
+
     fn evaluate_expr_compare(
         &mut self,
         expr_compare: &ast::ExprCompare,

crates/red_knot_python_semantic/src/types/property_tests.rs

@@ -75,6 +75,8 @@ fn arbitrary_core_type(g: &mut Gen) -> Ty {
         Ty::AbcClassLiteral("ABCMeta"),
         Ty::SubclassOfAbcClass("ABC"),
         Ty::SubclassOfAbcClass("ABCMeta"),
+        Ty::AlwaysTruthy,
+        Ty::AlwaysFalsy,
     ])
     .unwrap()
     .clone()