Repository: quora/pyanalyze
Pull request #491 (2022-03-05 20:32:10)

Fix type narrowing in else branch if constraint is stored in a variable


Patch:
diff --git a/docs/changelog.md b/docs/changelog.md
index a3594d3..67fd535 100644
--- a/docs/changelog.md
+++ b/docs/changelog.md
@@ -2,6 +2,8 @@
 
 ## Unreleased
 
+- Fix type narrowing in else branch if constraint is stored in a
+  variable (#491)
 - Fix incorrectly inferred `Never` return type for some function
   implementations (#490)
 - Infer precise call signatures for `TypedDict` types (#487)
diff --git a/pyanalyze/name_check_visitor.py b/pyanalyze/name_check_visitor.py
index 151d7b3..98b904d 100644
--- a/pyanalyze/name_check_visitor.py
+++ b/pyanalyze/name_check_visitor.py
@@ -98,6 +98,7 @@ from .stacked_scopes import (
     AbstractConstraint,
     Composite,
     CompositeIndex,
+    EquivalentConstraint,
     FunctionScope,
     Varname,
     Constraint,
@@ -3296,7 +3297,11 @@ class NameCheckVisitor(node_visitor.ReplacingNodeVisitor):
                 constraint = Constraint(
                     composite.varname, ConstraintType.is_truthy, True, None
                 )
-                return annotate_with_constraint(composite.value, constraint)
+                existing = extract_constraints(composite.value)
+                new_value, _ = unannotate_value(composite.value, ConstraintExtension)
+                return annotate_with_constraint(
+                    new_value, EquivalentConstraint.make([constraint, existing])
+                )
             else:
                 return composite.value
         else:
diff --git a/pyanalyze/stacked_scopes.py b/pyanalyze/stacked_scopes.py
index 8af61aa..898c69e 100644
--- a/pyanalyze/stacked_scopes.py
+++ b/pyanalyze/stacked_scopes.py
@@ -523,6 +523,42 @@ class PredicateProvider(AbstractConstraint):
         return NULL_CONSTRAINT
 
 
+@dataclass(frozen=True)
+class EquivalentConstraint(AbstractConstraint):
+    """Represents multiple constraints that are either all true or all false."""
+
+    constraints: Tuple[AbstractConstraint, ...]
+
+    def apply(self) -> Iterable["Constraint"]:
+        for cons in self.constraints:
+            yield from cons.apply()
+
+    def invert(self) -> "EquivalentConstraint":
+        # ~(A == B) -> ~A == ~B
+        return EquivalentConstraint(tuple([cons.invert() for cons in self.constraints]))
+
+    @classmethod
+    def make(cls, constraints: Iterable[AbstractConstraint]) -> AbstractConstraint:
+        processed = {}
+        for cons in constraints:
+            if isinstance(cons, EquivalentConstraint):
+                for subcons in cons.constraints:
+                    processed[id(subcons)] = subcons
+                continue
+            processed[id(cons)] = cons
+
+        final = list(processed.values())
+
+        if len(final) == 1:
+            (cons,) = final
+            return cons
+        return cls(tuple(final))
+
+    def __str__(self) -> str:
+        children = " == ".join(map(str, self.constraints))
+        return f"({children})"
+
+
 @dataclass(frozen=True)
 class AndConstraint(AbstractConstraint):
     """Represents the AND of two constraints."""


Test patch:
diff --git a/pyanalyze/test_stacked_scopes.py b/pyanalyze/test_stacked_scopes.py
index 3d8726b..00e630a 100644
--- a/pyanalyze/test_stacked_scopes.py
+++ b/pyanalyze/test_stacked_scopes.py
@@ -1775,4 +1775,17 @@ class TestInvalidation(TestNameCheckVisitorBase):
                     assert_is_value(key, TypedValue(int))
                     data = [ids, data[key]]
                 else:
+                    assert_is_value(key, KnownValue(None))
                     data = [ids]
+
+    @assert_passes()
+    def test_else(self) -> None:
+        from typing import Optional
+
+        def capybara(key: Optional[int]):
+            has_bias = key is not None
+            assert_is_value(key, TypedValue(int) | KnownValue(None))
+            if has_bias:
+                assert_is_value(key, TypedValue(int))
+            else:
+                assert_is_value(key, KnownValue(None))