Repository: quora/pyanalyze (python)

Issue:
walrus doesn't work in and/or
(Found by @GBeauregard on Discord)

Example test case:
```python

    @skip_before((3, 8))
    def test_and(self):
        self.assert_passes(
            """
            from typing import Optional

            def opt() -> Optional[int]:
                return None

            def capybara(cond):
                if (x := opt()) and cond:
                    assert_is_value(x, TypedValue(int))
                assert_is_value(x, TypedValue(int) | KnownValue(None))
            """
        )
        self.assert_passes(
            """
            def func(myvar: str, strset: set[str]):
                if (encoder_type := myvar) and myvar in strset:
                    print(encoder_type)
            """
        )
```

This is because we create a new subscope for and/or (so we can narrow things on the RHS) and throw away that subscope after the and/or finishes. But this needs to change so that name bindings from the walrus survive.

This probably also affects ternary: `(x := 2) if some_condition else (x := 1)`. That's much less likely to appear in real code though.

Pull request description:
Fix treatment of walrus in and/or and IfExp
Also ended up improving type inference for and/or in general.

Fixes #438


Patch:
diff --git a/docs/changelog.md b/docs/changelog.md
index ccde4c8..d3cd323 100644
--- a/docs/changelog.md
+++ b/docs/changelog.md
@@ -2,6 +2,8 @@
 
 ## Unreleased
 
+- Fix treatment of walrus operator in `and`, `or`, and `if/else`
+  expressions (#441)
 - Refactor `isinstance()` support (#440)
 - Exclude `Any[unreachable]` from unified values (#439)
 - Add support for `reveal_locals()` (#436)
diff --git a/pyanalyze/name_check_visitor.py b/pyanalyze/name_check_visitor.py
index 4cfd4d1..54a11c5 100644
--- a/pyanalyze/name_check_visitor.py
+++ b/pyanalyze/name_check_visitor.py
@@ -2590,22 +2590,34 @@ class NameCheckVisitor(node_visitor.ReplacingNodeVisitor):
         # we set a negative constraint.
 
         is_and = isinstance(node.op, ast.And)
+        stack = contextlib.ExitStack()
+        scopes = []
         out_constraints = []
-        with self.scopes.subscope():
-            values = []
-            left = node.values[:-1]
-            for condition in left:
-                new_value, constraint = self.constraint_from_condition(condition)
-                out_constraints.append(constraint)
+        values = []
+        constraint = NULL_CONSTRAINT
+        with stack:
+            for i, condition in enumerate(node.values):
+                is_last = i == len(node.values) - 1
+                scope = stack.enter_context(self.scopes.subscope())
+                scopes.append(scope)
                 if is_and:
                     self.add_constraint(condition, constraint)
-                    values.append(constrain_value(new_value, FALSY_CONSTRAINT))
                 else:
                     self.add_constraint(condition, constraint.invert())
+
+                new_value, constraint = self.constraint_from_condition(
+                    condition, check_boolability=not is_last
+                )
+                out_constraints.append(constraint)
+
+                if is_last:
+                    values.append(new_value)
+                elif is_and:
+                    values.append(constrain_value(new_value, FALSY_CONSTRAINT))
+                else:
                     values.append(constrain_value(new_value, TRUTHY_CONSTRAINT))
-            right_value = self._visit_possible_constraint(node.values[-1])
-            values.append(right_value)
-            out_constraints.append(extract_constraints(right_value))
+
+        self.scopes.combine_subscopes(scopes)
         constraint_cls = AndConstraint if is_and else OrConstraint
         constraint = constraint_cls.make(reversed(out_constraints))
         return annotate_with_constraint(unite_values(*values), constraint)
@@ -3433,12 +3445,13 @@ class NameCheckVisitor(node_visitor.ReplacingNodeVisitor):
 
     def visit_IfExp(self, node: ast.IfExp) -> Value:
         _, constraint = self.constraint_from_condition(node.test)
-        with self.scopes.subscope():
+        with self.scopes.subscope() as if_scope:
             self.add_constraint(node, constraint)
             then_val = self.visit(node.body)
-        with self.scopes.subscope():
+        with self.scopes.subscope() as else_scope:
             self.add_constraint(node, constraint.invert())
             else_val = self.visit(node.orelse)
+        self.scopes.combine_subscopes([if_scope, else_scope])
         return unite_values(then_val, else_val)
 
     def constraint_from_condition(


Test patch:
diff --git a/pyanalyze/test_name_check_visitor.py b/pyanalyze/test_name_check_visitor.py
index aad8c95..ca6a00d 100644
--- a/pyanalyze/test_name_check_visitor.py
+++ b/pyanalyze/test_name_check_visitor.py
@@ -2360,6 +2360,41 @@ class TestWalrus(TestNameCheckVisitorBase):
             """
         )
 
+    @skip_before((3, 8))
+    def test_and(self):
+        self.assert_passes(
+            """
+            from typing import Optional
+
+            def opt() -> Optional[int]:
+                return None
+
+            def capybara(cond):
+                if (x := opt()) and cond:
+                    assert_is_value(x, TypedValue(int))
+                assert_is_value(x, TypedValue(int) | KnownValue(None))
+            """
+        )
+        self.assert_passes(
+            """
+            from typing import Set
+
+            def func(myvar: str, strset: Set[str]) -> None:
+                if (encoder_type := myvar) and myvar in strset:
+                    print(encoder_type)
+            """
+        )
+
+    @skip_before((3, 8))
+    def test_if_exp(self):
+        self.assert_passes(
+            """
+            def capybara(cond):
+                (x := 2) if cond else (x := 1)
+                assert_is_value(x, KnownValue(2) | KnownValue(1))
+            """
+        )
+
     @skip_before((3, 8))
     def test_comprehension_scope(self):
         self.assert_passes(
diff --git a/pyanalyze/test_stacked_scopes.py b/pyanalyze/test_stacked_scopes.py
index 7e677a3..27e9d4f 100644
--- a/pyanalyze/test_stacked_scopes.py
+++ b/pyanalyze/test_stacked_scopes.py
@@ -1053,7 +1053,7 @@ class TestConstraints(TestNameCheckVisitorBase):
                 assert_is_value(y, KnownValue(True))
             else:
                 # no constraints from the inverse of an AND constraint
-                assert_is_value(x, AnyValue(AnySource.unannotated))
+                assert_is_value(x, KnownValue(True) | AnyValue(AnySource.unannotated))
                 assert_is_value(y, AnyValue(AnySource.unannotated))
 
         def kerodon(x):
@@ -1520,7 +1520,11 @@ class TestConstraints(TestNameCheckVisitorBase):
                     assert_is_value(
                         self.value,
                         MultiValuedValue(
-                            [AnyValue(AnySource.unannotated), KnownValue("")]
+                            [
+                                AnyValue(AnySource.unannotated),
+                                TypedValue(Foo),
+                                KnownValue(""),
+                            ]
                         ),
                     )
                 assert_is_value(