Repository: quora/pyanalyze
Pull request #492 (2022-03-05 21:27:44)

Enable exhaustive checking of enums and booleans


Patch:
diff --git a/docs/changelog.md b/docs/changelog.md
index 67fd535..8adcad8 100644
--- a/docs/changelog.md
+++ b/docs/changelog.md
@@ -2,6 +2,7 @@
 
 ## Unreleased
 
+- Enable exhaustive checking of enums and booleans (#492)
 - Fix type narrowing in else branch if constraint is stored in a
   variable (#491)
 - Fix incorrectly inferred `Never` return type for some function
diff --git a/pyanalyze/name_check_visitor.py b/pyanalyze/name_check_visitor.py
index 98b904d..1f73e66 100644
--- a/pyanalyze/name_check_visitor.py
+++ b/pyanalyze/name_check_visitor.py
@@ -84,6 +84,7 @@ from .options import (
     add_arguments,
 )
 from .patma import PatmaVisitor
+from .predicates import EqualsPredicate
 from .shared_options import Paths, ImportPaths, EnforceNoUnused
 from .reexport import ImplicitReexportTracker
 from .safe import (
@@ -2800,32 +2801,13 @@ class NameCheckVisitor(node_visitor.ReplacingNodeVisitor):
         if varname is None:
             return NULL_CONSTRAINT
         if isinstance(op, (ast.Is, ast.IsNot)):
+            predicate = EqualsPredicate(other_val, self, use_is=True)
             positive = isinstance(op, ast.Is)
-            return Constraint(varname, ConstraintType.is_value, positive, other_val)
+            return Constraint(varname, ConstraintType.predicate, positive, predicate)
         elif isinstance(op, (ast.Eq, ast.NotEq)):
-
-            def predicate_func(value: Value, positive: bool) -> Optional[Value]:
-                op = operator.eq if positive else operator.ne
-                if isinstance(value, KnownValue):
-                    try:
-                        result = op(value.val, other_val)
-                    except Exception:
-                        pass
-                    else:
-                        if not result:
-                            return None
-                elif positive:
-                    known_other = KnownValue(other_val)
-                    if value.is_assignable(known_other, self):
-                        return known_other
-                    else:
-                        return None
-                return value
-
+            predicate = EqualsPredicate(other_val, self)
             positive = isinstance(op, ast.Eq)
-            return Constraint(
-                varname, ConstraintType.predicate, positive, predicate_func
-            )
+            return Constraint(varname, ConstraintType.predicate, positive, predicate)
         elif isinstance(op, (ast.In, ast.NotIn)) and is_right:
 
             def predicate_func(value: Value, positive: bool) -> Optional[Value]:


Test patch:
diff --git a/pyanalyze/test_enum.py b/pyanalyze/test_enum.py
index bebc1e5..83ceb95 100644
--- a/pyanalyze/test_enum.py
+++ b/pyanalyze/test_enum.py
@@ -1,6 +1,6 @@
 # static analysis: ignore
 from .implementation import assert_is_value
-from .value import SubclassValue, TypedValue
+from .value import AnySource, AnyValue, KnownValue, SubclassValue, TypedValue
 from .test_name_check_visitor import TestNameCheckVisitorBase
 from .test_node_visitor import assert_passes
 
@@ -65,3 +65,31 @@ class TestEnum(TestNameCheckVisitorBase):
         class Foo(enum.Enum):
             a = 1
             b = 1  # E: duplicate_enum_member
+
+
+class TestNarrowing(TestNameCheckVisitorBase):
+    @assert_passes()
+    def test_exhaustive(self):
+        from enum import Enum
+
+        class X(Enum):
+            a = 1
+            b = 2
+
+        def capybara_eq(x: X):
+            if x == X.a:
+                assert_is_value(x, KnownValue(X.a))
+            else:
+                assert_is_value(x, KnownValue(X.b))
+
+        def capybara_is(x: X):
+            if x is X.a:
+                assert_is_value(x, KnownValue(X.a))
+            else:
+                assert_is_value(x, KnownValue(X.b))
+
+        def whatever(x):
+            if x == X.a:
+                assert_is_value(x, KnownValue(X.a))
+                return
+            assert_is_value(x, AnyValue(AnySource.unannotated))
diff --git a/pyanalyze/test_never.py b/pyanalyze/test_never.py
index 1cb8008..7f36a8c 100644
--- a/pyanalyze/test_never.py
+++ b/pyanalyze/test_never.py
@@ -6,7 +6,14 @@ from .test_name_check_visitor import TestNameCheckVisitorBase
 
 
 class TestAnnotations(TestNameCheckVisitorBase):
-    # TODO test typing_extensions.Never when it exists
+    @assert_passes()
+    def test_te_never(self):
+        from typing_extensions import Never
+
+        def capybara(n: Never, o: "Never"):
+            assert_is_value(n, NO_RETURN_VALUE)
+            assert_is_value(o, NO_RETURN_VALUE)
+
     @skip_before((3, 11))
     @assert_passes()
     def test_never(self):
diff --git a/pyanalyze/test_stacked_scopes.py b/pyanalyze/test_stacked_scopes.py
index 00e630a..ecb337e 100644
--- a/pyanalyze/test_stacked_scopes.py
+++ b/pyanalyze/test_stacked_scopes.py
@@ -821,6 +821,15 @@ class TestConstraints(TestNameCheckVisitorBase):
             assert y
             assert_is_value(y, KnownValue(True))
 
+    @assert_passes()
+    def test_bool_narrowing(self):
+        def capybara(x: bool):
+            assert_is_value(x, TypedValue(bool))
+            if x is True:
+                assert_is_value(x, KnownValue(True))
+            else:
+                assert_is_value(x, KnownValue(False))
+
     @assert_passes()
     def test_assert_falsy(self):
         def capybara(x):