[red knot] Fix narrowing for '… is not …' type guards, add '… is …' type guards

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

@@ -0,0 +1,26 @@
+# Narrowing for `is` conditionals
+
+## `is None`
+
+```py
+x = None if flag else 1
+
+if x is None:
+    # TODO the following should be simplified to 'None'
+    reveal_type(x)  # revealed: None | Literal[1] & None
+
+reveal_type(x)  # revealed: None | Literal[1]
+```
+
+## `is` for other types
+
+```py
+x = 345
+y = x if flag else None
+
+if y is x:
+    # TODO the following should be simplified to 'Literal[345]'
+    reveal_type(y)  # revealed: Literal[345] | None & Literal[345]
+
+reveal_type(y)  # revealed: Literal[345] | None
+```

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

@@ -0,0 +1,37 @@
+# Narrowing for `is not` conditionals
+
+## `is not None`
+
+The type guard removes `None` from the union type:
+
+```py
+x = None if flag else 1
+
+if x is not None:
+    reveal_type(x)  # revealed: Literal[1]
+
+reveal_type(x)  # revealed: None | Literal[1]
+```
+
+## `is not` for other singleton types
+
+```py
+x = True if flag else False
+reveal_type(x)  # revealed: bool
+
+if x is not False:
+    # TODO the following should be `Literal[True]`
+    reveal_type(x)  # revealed: bool & ~Literal[False]
+```
+
+## `is not` for non-singleton types
+
+Non-singleton types should *not* narrow the type:
+
+```py
+x = [1]
+y = [1]
+
+if x is not y:
+    reveal_type(x)  # revealed: list
+```

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

@@ -0,0 +1,17 @@
+# Narrowing for `match` statements
+
+## Single `match` pattern
+
+```py
+x = None if flag else 1
+reveal_type(x)  # revealed: None | Literal[1]
+
+y = 0
+
+match x:
+    case None:
+        y = x
+
+# TODO intersection simplification: should be just Literal[0] | None
+reveal_type(y)  # revealed: Literal[0] | None | Literal[1] & None
+```

crates/red_knot_python_semantic/src/types.rs

@@ -463,6 +463,52 @@ impl<'db> Type<'db> {
         self == other
     }
 
+    /// Return true if there is just a single inhabitant for this type.
+    ///
+    /// Note: This function aims to have no false positives, but might return `false`
+    /// for more complicated types that are actually singletons.
+    pub(crate) fn is_singleton(self, db: &'db dyn Db) -> bool {
+        match self {
+            Type::Any
+            | Type::Never
+            | Type::Unknown
+            | Type::Todo
+            | Type::Unbound
+            | Type::Module(..)
+            | Type::Instance(..)
+            | Type::IntLiteral(..)
+            | Type::StringLiteral(..)
+            | Type::BytesLiteral(..)
+            | Type::LiteralString => {
+                // Note: The literal types included in this pattern are not true singletons.
+                // There can be multiple Python objects (at different memory locations) that
+                // are both of type Literal[345], for example.
+                false
+            }
+            Type::None | Type::BooleanLiteral(_) | Type::Function(..) | Type::Class(..) => true,
+            Type::Tuple(tuple) => tuple.elements(db).iter().all(|ty| ty.is_singleton(db)),
+            Type::Union(..) => {
+                // There are some rare edge cases where a union type might be a singleton type.
+                // For example, a union with just one element (which itself is a singleton). Or
+                // a union with an empty type (e.g. Never | None). Here, we assume that such
+                // types would have been simplified to a different representation earlier and
+                // simply return false.
+                false
+            }
+            Type::Intersection(..) => {
+                // Intersection types are hard to analyze. The following types are technically
+                // all singleton types, but it is not straightforward to compute this. Again,
+                // we simply return false.
+                //
+                //   bool & ~Literal[False]`
+                //   None & (None | int)
+                //   (A | B) & (B | C)          with A, B, C disjunct and B a singleton
+                //
+                false
+            }
+        }
+    }
+
     /// Resolve a member access of a type.
     ///
     /// For example, if `foo` is `Type::Instance(<Bar>)`,
@@ -1514,6 +1560,7 @@ mod tests {
     enum Ty {
         Never,
         Unknown,
+        None,
         Any,
         IntLiteral(i64),
         BoolLiteral(bool),
@@ -1530,6 +1577,7 @@ mod tests {
             match self {
                 Ty::Never => Type::Never,
                 Ty::Unknown => Type::Unknown,
+                Ty::None => Type::None,
                 Ty::Any => Type::Any,
                 Ty::IntLiteral(n) => Type::IntLiteral(n),
                 Ty::StringLiteral(s) => {
@@ -1618,6 +1666,28 @@ mod tests {
         assert!(from.into_type(&db).is_equivalent_to(&db, to.into_type(&db)));
     }
 
+    #[test_case(Ty::None)]
+    #[test_case(Ty::BoolLiteral(true))]
+    #[test_case(Ty::BoolLiteral(false))]
+    #[test_case(Ty::Tuple(vec![]))]
+    #[test_case(Ty::Tuple(vec![Ty::None]))]
+    #[test_case(Ty::Tuple(vec![Ty::None, Ty::BoolLiteral(true)]))]
+    fn is_singleton(from: Ty) {
+        let db = setup_db();
+
+        assert!(from.into_type(&db).is_singleton(&db));
+    }
+
+    #[test_case(Ty::Never)]
+    #[test_case(Ty::IntLiteral(345))]
+    #[test_case(Ty::BuiltinInstance("str"))]
+    #[test_case(Ty::Union(vec![Ty::IntLiteral(1), Ty::IntLiteral(2)]))]
+    fn is_not_singleton(from: Ty) {
+        let db = setup_db();
+
+        assert!(!from.into_type(&db).is_singleton(&db));
+    }
+
     #[test_case(Ty::IntLiteral(1); "is_int_literal_truthy")]
     #[test_case(Ty::IntLiteral(-1))]
     #[test_case(Ty::StringLiteral("foo"))]

crates/red_knot_python_semantic/src/types/infer.rs

@@ -5421,53 +5421,6 @@ mod tests {
         Ok(())
     }
 
-    #[test]
-    fn narrow_not_none() -> anyhow::Result<()> {
-        let mut db = setup_db();
-
-        db.write_dedented(
-            "/src/a.py",
-            "
-            x = None if flag else 1
-            y = 0
-            if x is not None:
-                y = x
-            ",
-        )?;
-
-        assert_public_ty(&db, "/src/a.py", "x", "None | Literal[1]");
-        assert_public_ty(&db, "/src/a.py", "y", "Literal[0, 1]");
-
-        Ok(())
-    }
-
-    #[test]
-    fn narrow_singleton_pattern() {
-        let mut db = setup_db();
-
-        db.write_dedented(
-            "/src/a.py",
-            "
-            x = None if flag else 1
-            y = 0
-            match x:
-                case None:
-                    y = x
-            ",
-        )
-        .unwrap();
-
-        // TODO: The correct inferred type should be `Literal[0] | None` but currently the
-        // simplification logic doesn't account for this. The final type with parenthesis:
-        // `Literal[0] | None | (Literal[1] & None)`
-        assert_public_ty(
-            &db,
-            "/src/a.py",
-            "y",
-            "Literal[0] | None | Literal[1] & None",
-        );
-    }
-
     #[test]
     fn while_loop() -> anyhow::Result<()> {
         let mut db = setup_db();

crates/red_knot_python_semantic/src/types/narrow.rs

@@ -155,13 +155,20 @@ impl<'db> NarrowingConstraintsBuilder<'db> {
             let inference = infer_expression_types(self.db, expression);
             for (op, comparator) in std::iter::zip(&**ops, &**comparators) {
                 let comp_ty = inference.expression_ty(comparator.scoped_ast_id(self.db, scope));
-                if matches!(op, ast::CmpOp::IsNot) {
-                    let ty = IntersectionBuilder::new(self.db)
-                        .add_negative(comp_ty)
-                        .build();
-                    self.constraints.insert(symbol, ty);
-                };
-                // TODO other comparison types
+                match op {
+                    ast::CmpOp::IsNot if comp_ty.is_singleton(self.db) => {
+                        let ty = IntersectionBuilder::new(self.db)
+                            .add_negative(comp_ty)
+                            .build();
+                        self.constraints.insert(symbol, ty);
+                    }
+                    ast::CmpOp::Is => {
+                        self.constraints.insert(symbol, comp_ty);
+                    }
+                    _ => {
+                        // TODO other comparison types
+                    }
+                }
             }
         }
     }