Repository: quora/pyanalyze
Pull request #766 (2024-05-05 04:16:33)

Fix some pattern matching bugs


Patch:
diff --git a/docs/changelog.md b/docs/changelog.md
index 02ebc32..048533f 100644
--- a/docs/changelog.md
+++ b/docs/changelog.md
@@ -2,6 +2,8 @@
 
 ## Unreleased
 
+- Recognize exhaustive pattern matching (#766)
+- Narrow the types of variables assigned within complex patterns (#766)
 - New error code `generator_return` is raised when a generator does not
   return an iterable type, or an async generator does not return an async
   iterable type (#756)
diff --git a/pyanalyze/name_check_visitor.py b/pyanalyze/name_check_visitor.py
index 6ca7e6a..99627de 100644
--- a/pyanalyze/name_check_visitor.py
+++ b/pyanalyze/name_check_visitor.py
@@ -5649,10 +5649,20 @@ class NameCheckVisitor(node_visitor.ReplacingNodeVisitor):
 
                     self.yield_checker.reset_yield_checks()
 
-                with self.scopes.subscope() as else_scope:
-                    for constraint in constraints_to_apply:
-                        self.add_constraint(node, constraint)
-                    subscopes.append(else_scope)
+                self.match_subject = self.match_subject._replace(
+                    value=constrain_value(
+                        self.match_subject.value,
+                        AndConstraint.make(constraints_to_apply),
+                    )
+                )
+
+                if self.match_subject.value is NO_RETURN_VALUE:
+                    self._set_name_in_scope(LEAVES_SCOPE, node, NO_RETURN_VALUE)
+                else:
+                    with self.scopes.subscope() as else_scope:
+                        for constraint in constraints_to_apply:
+                            self.add_constraint(node, constraint)
+                        subscopes.append(else_scope)
                 self.scopes.combine_subscopes(subscopes)
 
     # Attribute checking
diff --git a/pyanalyze/patma.py b/pyanalyze/patma.py
index 04832ff..b8d19b5 100644
--- a/pyanalyze/patma.py
+++ b/pyanalyze/patma.py
@@ -177,6 +177,15 @@ class LenPredicate:
         return value
 
 
+@dataclass
+class AlwaysMatching:
+    def __call__(self, value: Value, positive: bool) -> Optional[Value]:
+        if positive:
+            return value
+        else:
+            return None
+
+
 @dataclass
 class PatmaVisitor(ast.NodeVisitor):
     visitor: "pyanalyze.name_check_visitor.NameCheckVisitor"
@@ -379,12 +388,14 @@ class PatmaVisitor(ast.NodeVisitor):
             self.visitor._set_name_in_scope(
                 node.name, node, self.visitor.match_subject.value
             )
-        return NULL_CONSTRAINT
+        return self.make_constraint(ConstraintType.predicate, AlwaysMatching())
 
     def visit_MatchAs(self, node: MatchAs) -> AbstractConstraint:
         val = self.visitor.match_subject.value
         if node.pattern is None:
-            constraint = NULL_CONSTRAINT
+            constraint = self.make_constraint(
+                ConstraintType.predicate, AlwaysMatching()
+            )
         else:
             constraint = self.visit(node.pattern)
 
@@ -405,12 +416,10 @@ class PatmaVisitor(ast.NodeVisitor):
         return OrConstraint.make(constraints)
 
     def generic_visit(self, node: ast.AST) -> AbstractConstraint:
-        return NULL_CONSTRAINT
+        raise NotImplementedError(f"Unsupported pattern node: {node}")
 
     def make_constraint(self, typ: ConstraintType, value: object) -> AbstractConstraint:
         varname = self.visitor.match_subject.varname
-        if varname is None:
-            return NULL_CONSTRAINT
         return Constraint(varname, typ, True, value)
 
     def check_impossible_pattern(self, node: ast.AST, value: Value) -> None:
diff --git a/pyanalyze/stacked_scopes.py b/pyanalyze/stacked_scopes.py
index 7efae94..da808c0 100644
--- a/pyanalyze/stacked_scopes.py
+++ b/pyanalyze/stacked_scopes.py
@@ -296,7 +296,7 @@ class Constraint(AbstractConstraint):
 
     """
 
-    varname: VarnameWithOrigin
+    varname: Optional[VarnameWithOrigin]
     """The :term:`varname` that the constraint applies to."""
     constraint_type: ConstraintType
     """Type of constraint. Determines the meaning of :attr:`value`."""
@@ -313,7 +313,9 @@ class Constraint(AbstractConstraint):
     )
 
     def __post_init__(self) -> None:
-        assert isinstance(self.varname, VarnameWithOrigin), self.varname
+        assert self.varname is None or isinstance(
+            self.varname, VarnameWithOrigin
+        ), self.varname
 
     def apply(self) -> Iterable["Constraint"]:
         yield self
@@ -1069,6 +1071,8 @@ class FunctionScope(Scope):
     def _add_single_constraint(
         self, constraint: Constraint, node: Node, state: VisitorState
     ) -> None:
+        if constraint.varname is None:
+            return
         for parent_varname, constraint_origin in constraint.varname.get_all_varnames():
             current_origin = self.get_origin(parent_varname, node, state)
             current_set = self._resolve_origin(current_origin)


Test patch:
diff --git a/pyanalyze/test_patma.py b/pyanalyze/test_patma.py
index e8d3b1f..f97481d 100644
--- a/pyanalyze/test_patma.py
+++ b/pyanalyze/test_patma.py
@@ -212,3 +212,47 @@ class TestPatma(TestNameCheckVisitorBase):
                         assert_is_value(p, KnownValue(Planet.earth))
             """
         )
+
+    @skip_before((3, 10))
+    def test_exhaustive(self):
+        self.assert_passes(
+            """
+            def f(x: object) -> int:
+                match x:
+                    case _:
+                        return 1
+
+            def g(x: object) -> int:  # E: missing_return
+                match x:
+                    case _ if x == 2:
+                        return 1
+
+            def some_func() -> object:
+                return 1
+
+            def h() -> int:
+                match some_func():
+                    case _:
+                        return 1
+
+            def i(x: bool) -> int:
+                match x:
+                    case True:
+                        return 1
+                    case False:
+                        return 2
+            """
+        )
+
+    @skip_before((3, 10))
+    def test_reassign_in_tuple(self):
+        self.assert_passes(
+            """
+            from typing_extensions import assert_type
+
+            def f(x: int | str) -> None:
+                match (x,):
+                    case (int() as x,):
+                        assert_type(x, int)
+            """
+        )