Propagate walrus narrowing from nested expressions

mypy/checker.py

@@ -148,11 +148,13 @@ def __init__(self) -> None:
     CallExpr,
     ClassDef,
     ComparisonExpr,
+    ConditionalExpr,
     Context,
     ContinueStmt,
     Decorator,
     DelStmt,
     DictExpr,
+    DictionaryComprehension,
     EllipsisExpr,
     Expression,
     ExpressionStmt,
@@ -161,6 +163,7 @@ def __init__(self) -> None:
     FuncBase,
     FuncDef,
     FuncItem,
+    GeneratorExpr,
     GlobalDecl,
     IfStmt,
     Import,
@@ -6593,6 +6596,38 @@ def find_isinstance_check(
         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 `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.
+        """
+        collector = WalrusAssignmentCollector()
+        node.accept(collector)
+        if not collector.assignments:
+            return
+
+        for type_map in (if_map, else_map):
+            narrowed = {literal_hash(expr) for expr in type_map}
+            for assignment in collector.assignments:
+                if literal_hash(assignment.target) in narrowed:
+                    # The condition narrows this target more precisely.
+                    continue
+                assigned_type = self.lookup_type_or_none(assignment.value)
+                if assigned_type is not None:
+                    type_map[assignment.target] = assigned_type
+
     def find_isinstance_check_helper(
         self, node: Expression, *, in_boolean_context: bool = True
     ) -> tuple[TypeMap, TypeMap]:
@@ -6709,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.
@@ -6722,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.
@@ -9781,6 +9818,46 @@ def collapse_walrus(e: Expression) -> Expression:
     return e
 
 
+class WalrusAssignmentCollector(TraverserVisitor):
+    """Collect the walrus assignments which an expression always performs.
+
+    Traversal stops at any expression which may leave its operands unevaluated,
+    or which evaluates them in a separate scope. Every collected assignment has
+    therefore taken place once the visited expression has been evaluated,
+    whatever value it produced.
+    """
+
+    def __init__(self) -> None:
+        self.assignments: list[AssignmentExpr] = []
+
+    def visit_assignment_expr(self, o: AssignmentExpr, /) -> None:
+        self.assignments.append(o)
+        o.value.accept(self)
+
+    def visit_op_expr(self, o: OpExpr, /) -> None:
+        if o.op in ("and", "or"):
+            # Short-circuiting; find_isinstance_check recurses into these itself.
+            return
+        super().visit_op_expr(o)
+
+    def visit_comparison_expr(self, o: ComparisonExpr, /) -> None:
+        # `a < b < c` stops at the first false comparison.
+        for operand in o.operands[:2]:
+            operand.accept(self)
+
+    def visit_conditional_expr(self, o: ConditionalExpr, /) -> None:
+        o.cond.accept(self)
+
+    def visit_generator_expr(self, o: GeneratorExpr, /) -> None:
+        """Skip generator expressions, and the comprehensions built on them."""
+
+    def visit_dictionary_comprehension(self, o: DictionaryComprehension, /) -> None:
+        """Skip dictionary comprehensions."""
+
+    def visit_lambda_expr(self, o: LambdaExpr, /) -> None:
+        """Skip lambdas, whose body is not evaluated here."""
+
+
 def find_last_var_assignment_line(n: Node, v: Var) -> int:
     """Find the highest line number of a potential assignment to variable within node.
 

test-data/unit/check-inference.test

@@ -4288,6 +4288,134 @@ def check_or_nested(maybe: bool) -> None:
         reveal_type(bar)  # N: Revealed type is "builtins.list[builtins.int]"
         reveal_type(baz)  # N: Revealed type is "builtins.list[builtins.int]"
 
+[case testInferWalrusAssignmentNestedInCondition]
+class Foo:
+    def __init__(self, value: bool) -> None:
+        self.value = value
+
+def truthy(x: object) -> bool: ...
+
+def check_binary_op(maybe: bool, n: int) -> None:
+    woo = None
+    if maybe and (woo := 5) + n:
+        reveal_type(woo)  # N: Revealed type is "builtins.int"
+    else:
+        reveal_type(woo)  # N: Revealed type is "builtins.int | None"
+
+def check_call(maybe: bool) -> None:
+    woo = None
+    if maybe and truthy(woo := Foo(True)):
+        reveal_type(woo)  # N: Revealed type is "__main__.Foo"
+    else:
+        reveal_type(woo)  # N: Revealed type is "__main__.Foo | None"
+
+def check_isinstance(maybe: bool) -> None:
+    woo = None
+    if maybe and isinstance(woo := Foo(True), Foo):
+        reveal_type(woo)  # N: Revealed type is "__main__.Foo"
+    else:
+        reveal_type(woo)  # N: Revealed type is "__main__.Foo | None"
+
+def check_comparison(maybe: bool, n: int) -> None:
+    woo = None
+    if maybe and (woo := 5) > n:
+        reveal_type(woo)  # N: Revealed type is "builtins.int"
+    else:
+        reveal_type(woo)  # N: Revealed type is "builtins.int | None"
+
+def check_unary_op(maybe: bool) -> None:
+    woo = None
+    if maybe and -(woo := 5):
+        reveal_type(woo)  # N: Revealed type is "builtins.int"
+    else:
+        reveal_type(woo)  # N: Revealed type is "builtins.int | None"
+
+def check_or(maybe: bool, n: int) -> None:
+    woo = None
+    if maybe or (woo := 5) + n:
+        reveal_type(woo)  # N: Revealed type is "builtins.int | None"
+    else:
+        reveal_type(woo)  # N: Revealed type is "builtins.int"
+
+def check_nested_walrus(maybe: bool, n: int) -> None:
+    foo = None
+    bar = None
+    if maybe and (foo := (bar := 5)) + n:
+        reveal_type(foo)  # N: Revealed type is "builtins.int"
+        reveal_type(bar)  # N: Revealed type is "builtins.int"
+    else:
+        reveal_type(foo)  # N: Revealed type is "builtins.int | None"
+        reveal_type(bar)  # N: Revealed type is "builtins.int | None"
+
+def check_nested_and(maybe: bool) -> None:
+    # Nested short-circuiting operators are not walked into, but they are still
+    # handled, because find_isinstance_check recurses into them itself.
+    woo = None
+    if maybe and (1 and (woo := 5)):
+        reveal_type(woo)  # N: Revealed type is "builtins.int"
+[builtins fixtures/len.pyi]
+
+[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 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 maybe and (1 if (woo := 5) else 0):
+        reveal_type(woo)  # N: Revealed type is "builtins.int"
+
+def check_comprehension(maybe: bool, xs: List[int]) -> None:
+    woo = None
+    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(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 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
+
+# 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(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(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]
 from typing import Any, Optional, TypeVar