From 2296627528e94a8ebc55e958598b36df811952fc Mon Sep 17 00:00:00 2001
From: David Peter <sharkdp@users.noreply.github.com>
Date: Tue, 5 Nov 2024 19:48:52 +0100
Subject: [PATCH] [red-knot] Precise inference for identity checks (#14109)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

## Summary

Adds more precise type inference for `… is …` and `… is not …` identity
checks in some limited cases where we statically know the answer to be
either `Literal[True]` or `Literal[False]`.

I found this helpful while working on type inference for comparisons
involving intersection types, but I'm not sure if this is at all useful
for real world code (where the answer is most probably *not* statically
known). Note that we already have *type narrowing* for identity tests.
So while we are already able to generate constraints for things like `if
x is None`, we can now — in some limited cases — make an even stronger
conclusion and infer that the test expression itself is `Literal[False]`
(branch never taken) or `Literal[True]` (branch always taken).

## Test Plan

New Markdown tests
---
 .../mdtest/comparison/identity_tests.md       | 40 +++++++++++++++++++
 .../src/types/infer.rs                        | 24 ++++++++++-
 2 files changed, 62 insertions(+), 2 deletions(-)
 create mode 100644 crates/red_knot_python_semantic/resources/mdtest/comparison/identity_tests.md

diff --git a/crates/red_knot_python_semantic/resources/mdtest/comparison/identity_tests.md b/crates/red_knot_python_semantic/resources/mdtest/comparison/identity_tests.md
new file mode 100644
index 0000000000..bf162efa0a
--- /dev/null
+++ b/crates/red_knot_python_semantic/resources/mdtest/comparison/identity_tests.md
@@ -0,0 +1,40 @@
+# Identity tests
+
+```py
+class A: ...
+
+def get_a() -> A: ...
+def get_object() -> object: ...
+
+a1 = get_a()
+a2 = get_a()
+
+n1 = None
+n2 = None
+
+o = get_object()
+
+reveal_type(a1 is a1)  # revealed: bool
+reveal_type(a1 is a2)  # revealed: bool
+
+reveal_type(n1 is n1)  # revealed: Literal[True]
+reveal_type(n1 is n2)  # revealed: Literal[True]
+
+reveal_type(a1 is n1)  # revealed: Literal[False]
+reveal_type(n1 is a1)  # revealed: Literal[False]
+
+reveal_type(a1 is o)  # revealed: bool
+reveal_type(n1 is o)  # revealed: bool
+
+reveal_type(a1 is not a1)  # revealed: bool
+reveal_type(a1 is not a2)  # revealed: bool
+
+reveal_type(n1 is not n1)  # revealed: Literal[False]
+reveal_type(n1 is not n2)  # revealed: Literal[False]
+
+reveal_type(a1 is not n1)  # revealed: Literal[True]
+reveal_type(n1 is not a1)  # revealed: Literal[True]
+
+reveal_type(a1 is not o)  # revealed: bool
+reveal_type(n1 is not o)  # revealed: bool
+```
diff --git a/crates/red_knot_python_semantic/src/types/infer.rs b/crates/red_knot_python_semantic/src/types/infer.rs
index fd4b21758b..c4e0594a1f 100644
--- a/crates/red_knot_python_semantic/src/types/infer.rs
+++ b/crates/red_knot_python_semantic/src/types/infer.rs
@@ -3388,8 +3388,28 @@ impl<'db> TypeInferenceBuilder<'db> {
                     ast::CmpOp::NotIn => {
                         membership_test_comparison(MembershipTestCompareOperator::NotIn)
                     }
-                    ast::CmpOp::Is => Ok(KnownClass::Bool.to_instance(self.db)),
-                    ast::CmpOp::IsNot => Ok(KnownClass::Bool.to_instance(self.db)),
+                    ast::CmpOp::Is => {
+                        if left.is_disjoint_from(self.db, right) {
+                            Ok(Type::BooleanLiteral(false))
+                        } else if left.is_singleton(self.db)
+                            && left.is_equivalent_to(self.db, right)
+                        {
+                            Ok(Type::BooleanLiteral(true))
+                        } else {
+                            Ok(KnownClass::Bool.to_instance(self.db))
+                        }
+                    }
+                    ast::CmpOp::IsNot => {
+                        if left.is_disjoint_from(self.db, right) {
+                            Ok(Type::BooleanLiteral(true))
+                        } else if left.is_singleton(self.db)
+                            && left.is_equivalent_to(self.db, right)
+                        {
+                            Ok(Type::BooleanLiteral(false))
+                        } else {
+                            Ok(KnownClass::Bool.to_instance(self.db))
+                        }
+                    }
                 }
             }
             // TODO: handle more types