diff --git a/mypy/checker.py b/mypy/checker.py index 813939ca4964..ddde450baa9d 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, @@ -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. diff --git a/test-data/unit/check-inference.test b/test-data/unit/check-inference.test index a5b3ae7238a5..89768fd4eaa7 100644 --- a/test-data/unit/check-inference.test +++ b/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