Repository: quora/pyanalyze
Pull request #387 (2022-01-09 18:13:06)

Support type narrowing on enums and bools in match statements


Patch:
diff --git a/docs/changelog.md b/docs/changelog.md
index ed1038e..4f88906 100644
--- a/docs/changelog.md
+++ b/docs/changelog.md
@@ -4,6 +4,8 @@
 
 - Do not error on boolean operations on values typed
   as `object` (#388)
+- Support type narrowing on enum types and `bool`
+  in `match` statements (#387)
 - Support some imports from stub-only modules (#386)
 - Support type evaluation functions in stubs (#386)
 - Support `TypedDict` in stubs (#386)
diff --git a/pyanalyze/__init__.py b/pyanalyze/__init__.py
index 042da44..091a7d4 100644
--- a/pyanalyze/__init__.py
+++ b/pyanalyze/__init__.py
@@ -25,6 +25,7 @@ from . import method_return_type
 from . import node_visitor
 from . import options
 from . import patma
+from . import predicates
 from . import reexport
 from . import safe
 from . import shared_options
@@ -60,3 +61,4 @@ used(suggested_type)
 used(options)
 used(shared_options)
 used(functions)
+used(predicates)
diff --git a/pyanalyze/name_check_visitor.py b/pyanalyze/name_check_visitor.py
index fafdb03..dee5410 100644
--- a/pyanalyze/name_check_visitor.py
+++ b/pyanalyze/name_check_visitor.py
@@ -4479,6 +4479,12 @@ class NameCheckVisitor(node_visitor.ReplacingNodeVisitor):
                 with self.scopes.subscope() as case_scope:
                     for constraint in constraints_to_apply:
                         self.add_constraint(case, constraint)
+                    self.match_subject = self.match_subject._replace(
+                        value=constrain_value(
+                            self.match_subject.value,
+                            AndConstraint.make(constraints_to_apply),
+                        )
+                    )
 
                     pattern_constraint = patma_visitor.visit(case.pattern)
                     constraints = [pattern_constraint]
diff --git a/pyanalyze/patma.py b/pyanalyze/patma.py
index 9a0ad86..f2cbeb3 100644
--- a/pyanalyze/patma.py
+++ b/pyanalyze/patma.py
@@ -29,6 +29,7 @@ from .signature import MappingValue
 from .annotations import type_from_value
 from .extensions import CustomCheck
 from .error_code import ErrorCode
+from .predicates import EqualsPredicate, IsAssignablePredicate
 from .stacked_scopes import (
     NULL_CONSTRAINT,
     AbstractConstraint,
@@ -38,7 +39,6 @@ from .stacked_scopes import (
     ConstraintType,
     OrConstraint,
     constrain_value,
-    IsAssignablePredicate,
 )
 from .value import (
     AnnotatedValue,
@@ -178,7 +178,10 @@ class PatmaVisitor(ast.NodeVisitor):
 
     def visit_MatchSingleton(self, node: MatchSingleton) -> AbstractConstraint:
         self.check_impossible_pattern(node, KnownValue(node.value))
-        return self.make_constraint(ConstraintType.is_value, node.value)
+        return self.make_constraint(
+            ConstraintType.predicate,
+            EqualsPredicate(node.value, self.visitor, use_is=True),
+        )
 
     def visit_MatchValue(self, node: MatchValue) -> AbstractConstraint:
         pattern_val = self.visitor.visit(node.value)
@@ -191,26 +194,9 @@ class PatmaVisitor(ast.NodeVisitor):
             )
             return NULL_CONSTRAINT
 
-        def predicate_func(value: Value, positive: bool) -> Optional[Value]:
-            if isinstance(value, KnownValue):
-                try:
-                    if positive:
-                        result = value.val == pattern_val.val
-                    else:
-                        result = value.val != pattern_val.val
-                except Exception:
-                    pass
-                else:
-                    if not result:
-                        return None
-            elif positive:
-                if value.is_assignable(pattern_val, self.visitor):
-                    return pattern_val
-                else:
-                    return None
-            return value
-
-        return self.make_constraint(ConstraintType.predicate, predicate_func)
+        return self.make_constraint(
+            ConstraintType.predicate, EqualsPredicate(pattern_val.val, self.visitor)
+        )
 
     def visit_MatchSequence(self, node: MatchSequence) -> AbstractConstraint:
         self.check_impossible_pattern(node, MatchableSequence)
diff --git a/pyanalyze/predicates.py b/pyanalyze/predicates.py
new file mode 100644
index 0000000..bccf13b
--- /dev/null
+++ b/pyanalyze/predicates.py
@@ -0,0 +1,95 @@
+"""
+
+Reusable predicates.
+
+"""
+from dataclasses import dataclass
+import enum
+import operator
+from typing import Optional
+
+from .safe import safe_issubclass
+from .value import (
+    KnownValue,
+    TypedValue,
+    Value,
+    unite_values,
+    CanAssignContext,
+    is_overlapping,
+    unannotate,
+)
+
+
+@dataclass
+class IsAssignablePredicate:
+    """Predicate that filters out values that are not assignable to pattern_value."""
+
+    pattern_value: Value
+    ctx: CanAssignContext
+    positive_only: bool
+
+    def __call__(self, value: Value, positive: bool) -> Optional[Value]:
+        compatible = is_overlapping(self.pattern_value, value, self.ctx)
+        if positive:
+            if not compatible:
+                return None
+            if value.is_assignable(self.pattern_value, self.ctx):
+                return self.pattern_value
+        elif not self.positive_only:
+            if compatible:
+                return None
+        return value
+
+
+_OPERATOR = {
+    (True, True): operator.is_,
+    (False, True): operator.is_not,
+    (True, False): operator.eq,
+    (False, False): operator.ne,
+}
+
+
+@dataclass
+class EqualsPredicate:
+    """Predicate that filters out values that are not equal to pattern_val."""
+
+    pattern_val: object
+    ctx: CanAssignContext
+    use_is: bool = False
+
+    def __call__(self, value: Value, positive: bool) -> Optional[Value]:
+        inner_value = unannotate(value)
+        if isinstance(inner_value, KnownValue):
+            op = _OPERATOR[(positive, self.use_is)]
+            try:
+                result = op(inner_value.val, self.pattern_val)
+            except Exception:
+                pass
+            else:
+                if not result:
+                    return None
+        elif positive:
+            known_self = KnownValue(self.pattern_val)
+            if value.is_assignable(known_self, self.ctx):
+                return known_self
+            else:
+                return None
+        else:
+            pattern_type = type(self.pattern_val)
+            if pattern_type is bool:
+                simplified = unannotate(value)
+                if isinstance(simplified, TypedValue) and simplified.typ is bool:
+                    return KnownValue(not self.pattern_val)
+            elif safe_issubclass(pattern_type, enum.Enum):
+                simplified = unannotate(value)
+                if isinstance(simplified, TypedValue) and simplified.typ is type(
+                    self.pattern_val
+                ):
+                    return unite_values(
+                        *[
+                            KnownValue(val)
+                            for val in pattern_type
+                            if val is not self.pattern_val
+                        ]
+                    )
+        return value
diff --git a/pyanalyze/stacked_scopes.py b/pyanalyze/stacked_scopes.py
index 9ee78ee..7aa456c 100644
--- a/pyanalyze/stacked_scopes.py
+++ b/pyanalyze/stacked_scopes.py
@@ -67,8 +67,6 @@ from .value import (
     unite_and_simplify,
     unite_values,
     flatten_values,
-    CanAssignContext,
-    is_overlapping,
 )
 
 T = TypeVar("T")
@@ -516,25 +514,6 @@ class PredicateProvider(AbstractConstraint):
         return NULL_CONSTRAINT
 
 
-@dataclass
-class IsAssignablePredicate:
-    pattern_value: Value
-    ctx: CanAssignContext
-    positive_only: bool
-
-    def __call__(self, value: Value, positive: bool) -> Optional[Value]:
-        compatible = is_overlapping(self.pattern_value, value, self.ctx)
-        if positive:
-            if not compatible:
-                return None
-            if value.is_assignable(self.pattern_value, self.ctx):
-                return self.pattern_value
-        elif not self.positive_only:
-            if compatible:
-                return None
-        return value
-
-
 @dataclass(frozen=True)
 class AndConstraint(AbstractConstraint):
     """Represents the AND of two constraints."""
diff --git a/pyanalyze/type_evaluation.py b/pyanalyze/type_evaluation.py
index dfd043a..4d5348d 100644
--- a/pyanalyze/type_evaluation.py
+++ b/pyanalyze/type_evaluation.py
@@ -26,12 +26,12 @@ from typing import (
     Union,
 )
 
+from .predicates import IsAssignablePredicate
 from .stacked_scopes import (
     Constraint,
     ConstraintType,
     VarnameWithOrigin,
     constrain_value,
-    IsAssignablePredicate,
 )
 from .safe import all_of_type
 from .value import (


Test patch:
diff --git a/pyanalyze/test_boolability.py b/pyanalyze/test_boolability.py
index 7c7be40..53639b4 100644
--- a/pyanalyze/test_boolability.py
+++ b/pyanalyze/test_boolability.py
@@ -88,6 +88,7 @@ def test_get_boolability() -> None:
     assert Boolability.erroring_bool == get_boolability(future)
     assert Boolability.boolable == get_boolability(TypedValue(object))
     assert Boolability.boolable == get_boolability(TypedValue(int))
+    assert Boolability.boolable == get_boolability(TypedValue(object))
 
     # MultiValuedValue and AnnotatedValue
     assert Boolability.erroring_bool == get_boolability(
diff --git a/pyanalyze/test_patma.py b/pyanalyze/test_patma.py
index 19b675b..25f1128 100644
--- a/pyanalyze/test_patma.py
+++ b/pyanalyze/test_patma.py
@@ -158,3 +158,55 @@ class TestPatma(TestNameCheckVisitorBase):
                         pass
             """
         )
+
+    @skip_before((3, 10))
+    def test_bool_narrowing(self):
+        self.assert_passes(
+            """
+            class X:
+                true = True
+
+            def capybara(b: bool):
+                match b:
+                    # Make sure we hit the MatchValue case, not MatchSingleton
+                    case X.true:
+                        assert_is_value(b, KnownValue(True))
+                    case _ as b2:
+                        assert_is_value(b, KnownValue(False))
+                        assert_is_value(b2, KnownValue(False))
+            """
+        )
+        self.assert_passes(
+            """
+            def capybara(b: bool):
+                match b:
+                    case True:
+                        assert_is_value(b, KnownValue(True))
+                    case _ as b2:
+                        assert_is_value(b, KnownValue(False))
+                        assert_is_value(b2, KnownValue(False))
+            """
+        )
+
+    @skip_before((3, 10))
+    def test_enum_narrowing(self):
+        self.assert_passes(
+            """
+            from enum import Enum
+
+            class Planet(Enum):
+                mercury = 1
+                venus = 2
+                earth = 3
+
+            def capybara(p: Planet):
+                match p:
+                    case Planet.mercury:
+                        assert_is_value(p, KnownValue(Planet.mercury))
+                    case Planet.venus:
+                        assert_is_value(p, KnownValue(Planet.venus))
+                    case _ as p2:
+                        assert_is_value(p2, KnownValue(Planet.earth))
+                        assert_is_value(p, KnownValue(Planet.earth))
+            """
+        )