Only propagate walrus narrowing where the binder cannot carry it

Author: Endika

mypy/checker.py

@@ -6592,27 +6592,27 @@ def find_isinstance_check(
         if_map, else_map = self.find_isinstance_check_helper(
             node, in_boolean_context=in_boolean_context
         )
-        self.propagate_walrus_assignments(node, if_map, else_map)
         new_if_map = self.propagate_up_typemap_info(if_map)
         new_else_map = self.propagate_up_typemap_info(else_map)
         return new_if_map, new_else_map
 
     def propagate_walrus_assignments(
         self, node: Expression, if_map: TypeMap, else_map: TypeMap
     ) -> None:
-        """Narrow the targets of walrus assignments nested within a condition.
-
-        Such an assignment has already happened by the time the condition has
-        been evaluated, so the assigned type applies to both branches, and both
-        maps are updated in place. Which branches are actually reached is decided
-        by the callers combining these maps: `and` carries the right operand's if
-        map into the if branch, and `or` carries its else map into the else
-        branch.
+        """Narrow the targets of walrus assignments nested within `node`.
+
+        Only used for the right operand of `and` and `or`. Elsewhere the operand
+        is always evaluated, so the binder already carries the assignment and
+        adding it here would only widen the result: an entry makes the branches
+        join through the target's declaration, which may be wider than the type
+        assigned.
+
+        The assignment has happened once `node` has been evaluated, whatever
+        value it produced, so both maps are updated in place. Which branch that
+        reaches is decided by the caller combining them: `and` carries the right
+        operand's if map into the if branch, and `or` carries its else map into
+        the else branch.
         """
-        if isinstance(node, NameExpr):
-            # The most common condition, and it has no subexpressions.
-            return
-
         collector = WalrusAssignmentCollector()
         node.accept(collector)
         if not collector.assignments:
@@ -6744,6 +6744,7 @@ def find_isinstance_check_helper(
         elif isinstance(node, OpExpr) and node.op == "and":
             left_if_vars, left_else_vars = self.find_isinstance_check(node.left)
             right_if_vars, right_else_vars = self.find_isinstance_check(node.right)
+            self.propagate_walrus_assignments(node.right, right_if_vars, right_else_vars)
 
             # (e1 and e2) is true if both e1 and e2 are true,
             # and false if at least one of e1 and e2 is false.
@@ -6757,6 +6758,7 @@ def find_isinstance_check_helper(
         elif isinstance(node, OpExpr) and node.op == "or":
             left_if_vars, left_else_vars = self.find_isinstance_check(node.left)
             right_if_vars, right_else_vars = self.find_isinstance_check(node.right)
+            self.propagate_walrus_assignments(node.right, right_if_vars, right_else_vars)
 
             # (e1 or e2) is true if at least one of e1 or e2 is true,
             # and false if both e1 and e2 are false.

test-data/unit/check-inference.test

@@ -4358,48 +4358,62 @@ def check_nested_and(maybe: bool) -> None:
 [case testInferWalrusAssignmentNestedInConditionNotAlwaysEvaluated]
 from typing import List
 
+# Each condition puts the walrus on the right of an `and`, which is where the
+# assignment is not carried by the binder and this narrowing applies.
+
 def check_ternary_branch(maybe: bool) -> None:
     woo = None
-    if 1 if maybe else (woo := 5):
+    if maybe and (1 if maybe else (woo := 5)):
         reveal_type(woo)  # N: Revealed type is "builtins.int | None"
     else:
         reveal_type(woo)  # N: Revealed type is "builtins.int | None"
 
 def check_ternary_condition(maybe: bool) -> None:
     woo = None
-    if 1 if (woo := 5) else 0:
+    if maybe and (1 if (woo := 5) else 0):
         reveal_type(woo)  # N: Revealed type is "builtins.int"
 
-def check_comprehension(xs: List[int]) -> None:
+def check_comprehension(maybe: bool, xs: List[int]) -> None:
     woo = None
-    if [y for y in xs if (woo := y)]:
+    if maybe and [y for y in xs if (woo := y)]:
         reveal_type(woo)  # N: Revealed type is "builtins.int | None"
 
-def check_chained_comparison(a: int, b: int) -> None:
-    # The else branch should stay optional, and does not. Pre-existing: chain
-    # operands are checked in one binder frame, so the assignment is recorded
-    # even when short-circuiting means it never ran.
+def check_chained_comparison(maybe: bool, a: int, b: int) -> None:
+    # Conservative: entering the branch does imply a < b was true and so the
+    # walrus ran, but operands after the second are not walked into.
     woo = None
-    if a < b < (woo := 5):
-        reveal_type(woo)  # N: Revealed type is "builtins.int"
-    else:
-        reveal_type(woo)  # N: Revealed type is "builtins.int"
+    if maybe and a < b < (woo := 5):
+        reveal_type(woo)  # N: Revealed type is "builtins.int | None"
 [builtins fixtures/len.pyi]
 
 [case testInferWalrusAssignmentDoesNotWeakenNarrowing]
 from typing import Optional, Union
 
-def check_truthiness(val: Optional[int]) -> None:
-    if x := val:
+# The walrus goes on the right of an `and` so that the narrowing added for the
+# assignment has to give way to the more precise narrowing from the condition.
+
+def check_truthiness(maybe: bool, val: Optional[int]) -> None:
+    if maybe and (x := val):
         reveal_type(x)  # N: Revealed type is "builtins.int"
 
-def check_isinstance(val: Union[int, str]) -> None:
-    if isinstance(x := val, int):
+def check_isinstance(maybe: bool, val: Union[int, str]) -> None:
+    if maybe and isinstance(x := val, int):
         reveal_type(x)  # N: Revealed type is "builtins.int"
 
-def check_is_not_none(val: Optional[int]) -> None:
-    if (x := val) is not None:
+def check_is_not_none(maybe: bool, val: Optional[int]) -> None:
+    if maybe and (x := val) is not None:
         reveal_type(x)  # N: Revealed type is "builtins.int"
+
+def truthy(x: object) -> bool: ...
+
+def check_declaration_wider_than_assignment(val: Optional[int], n: int) -> None:
+    # An operand that is always evaluated must not be given a map entry: the
+    # branches would then join through the declaration of x, which is wider than
+    # what the walrus assigned. Reported by mypy_primer against rotki.
+    x = val
+    if truthy(x := n):
+        pass
+    reveal_type(x)  # N: Revealed type is "builtins.int"
 [builtins fixtures/isinstancelist.pyi]
 
 [case testInferOptionalAgainstAny]