From 9315c0ba71f9e738ced42588680aed36349b0c51 Mon Sep 17 00:00:00 2001 From: Endika Iglesias Date: Mon, 3 Aug 2026 12:04:39 +0200 Subject: [PATCH 1/3] Propagate walrus narrowing from nested expressions --- mypy/checker.py | 75 ++++++++++++++++++ test-data/unit/check-inference.test | 113 ++++++++++++++++++++++++++++ 2 files changed, 188 insertions(+) diff --git a/mypy/checker.py b/mypy/checker.py index 813939ca49646..84ccdf1894065 100644 --- a/mypy/checker.py +++ b/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, @@ -6589,10 +6592,42 @@ 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. + """ + if isinstance(node, NameExpr): + # The most common condition, and it has no subexpressions. + return + + 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]: @@ -9781,6 +9816,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. diff --git a/test-data/unit/check-inference.test b/test-data/unit/check-inference.test index a5b3ae7238a5a..e917ae08178f3 100644 --- a/test-data/unit/check-inference.test +++ b/test-data/unit/check-inference.test @@ -4288,6 +4288,119 @@ 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 + +def check_ternary_branch(maybe: bool) -> None: + woo = None + if 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: + reveal_type(woo) # N: Revealed type is "builtins.int" + +def check_comprehension(xs: List[int]) -> None: + woo = None + if [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: the + # operand is narrowed through collapse_walrus in comparison_type_narrowing_helper. + 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" +[builtins fixtures/len.pyi] + +[case testInferWalrusAssignmentDoesNotWeakenNarrowing] +from typing import Optional, Union + +def check_truthiness(val: Optional[int]) -> None: + if x := val: + reveal_type(x) # N: Revealed type is "builtins.int" + +def check_isinstance(val: Union[int, str]) -> None: + if 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: + reveal_type(x) # N: Revealed type is "builtins.int" +[builtins fixtures/isinstancelist.pyi] + [case testInferOptionalAgainstAny] from typing import Any, Optional, TypeVar From 7c7263648d28a44a77c2ae65d6cf589c3a787e73 Mon Sep 17 00:00:00 2001 From: Endika Iglesias Date: Mon, 3 Aug 2026 12:47:33 +0200 Subject: [PATCH 2/3] Correct the root cause noted for the pinned chained comparison --- test-data/unit/check-inference.test | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/test-data/unit/check-inference.test b/test-data/unit/check-inference.test index e917ae08178f3..05408f5de9917 100644 --- a/test-data/unit/check-inference.test +++ b/test-data/unit/check-inference.test @@ -4376,8 +4376,9 @@ def check_comprehension(xs: List[int]) -> None: 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: the - # operand is narrowed through collapse_walrus in comparison_type_narrowing_helper. + # 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. woo = None if a < b < (woo := 5): reveal_type(woo) # N: Revealed type is "builtins.int" From fa6ef093d964d752a1fedf0d15d21381ef9773e6 Mon Sep 17 00:00:00 2001 From: Endika Iglesias Date: Mon, 3 Aug 2026 13:38:51 +0200 Subject: [PATCH 3/3] Only propagate walrus narrowing where the binder cannot carry it --- mypy/checker.py | 28 ++++++++-------- test-data/unit/check-inference.test | 50 ++++++++++++++++++----------- 2 files changed, 47 insertions(+), 31 deletions(-) diff --git a/mypy/checker.py b/mypy/checker.py index 84ccdf1894065..ddde450baa9de 100644 --- a/mypy/checker.py +++ b/mypy/checker.py @@ -6592,7 +6592,6 @@ 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 @@ -6600,19 +6599,20 @@ def find_isinstance_check( 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. diff --git a/test-data/unit/check-inference.test b/test-data/unit/check-inference.test index 05408f5de9917..89768fd4eaa70 100644 --- a/test-data/unit/check-inference.test +++ b/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]