Skip to content

Commit 12166a7

Browse files
committed
Fix narrowing information not propagated in assignment and boolean expressions
Fixes #8925. Consider the following narrowing (more realistically with TypedDicts), and the current outcome: ```py class A: tag: Literal["A"] class B: tag: Literal["B"] abo: A | B | None if abo is not None and abo.tag == "A": reveal_type(abo.tag) # Type is Literal["A"] reveal_type(abo) # Type is A | B ``` The RHS of the comparison correctly takes into account the LHS, and the `abo.tag` expression is correctly narrowed based on it, but this does not propagate upward to then narrow `abo` to `A`. The problem is that `and`/`or/`not`/assignment expressions recurse using the `find_isinstance_check_helper` function, which omits the `propagate_up_typemap_info` calls that its parent function `find_isinstance_check` performs. Fix this by replacing these recursive `find_isinstance_check_helper` calls with `find_isinstance_check`. This might not always be necessary, but I do not think that always propagating upwards should do any harm, anyways.
1 parent 8e82171 commit 12166a7

File tree

2 files changed

+27
-7
lines changed

2 files changed

+27
-7
lines changed

mypy/checker.py

Lines changed: 7 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -4473,14 +4473,14 @@ def has_no_custom_eq_checks(t: Type) -> bool:
44734473
if_map = {}
44744474
else_map = {}
44754475

4476-
if_assignment_map, else_assignment_map = self.find_isinstance_check_helper(node.target)
4476+
if_assignment_map, else_assignment_map = self.find_isinstance_check(node.target)
44774477

44784478
if if_assignment_map is not None:
44794479
if_map.update(if_assignment_map)
44804480
if else_assignment_map is not None:
44814481
else_map.update(else_assignment_map)
44824482

4483-
if_condition_map, else_condition_map = self.find_isinstance_check_helper(node.value)
4483+
if_condition_map, else_condition_map = self.find_isinstance_check(node.value)
44844484

44854485
if if_condition_map is not None:
44864486
if_map.update(if_condition_map)
@@ -4492,23 +4492,23 @@ def has_no_custom_eq_checks(t: Type) -> bool:
44924492
(None if else_assignment_map is None or else_condition_map is None else else_map),
44934493
)
44944494
elif isinstance(node, OpExpr) and node.op == 'and':
4495-
left_if_vars, left_else_vars = self.find_isinstance_check_helper(node.left)
4496-
right_if_vars, right_else_vars = self.find_isinstance_check_helper(node.right)
4495+
left_if_vars, left_else_vars = self.find_isinstance_check(node.left)
4496+
right_if_vars, right_else_vars = self.find_isinstance_check(node.right)
44974497

44984498
# (e1 and e2) is true if both e1 and e2 are true,
44994499
# and false if at least one of e1 and e2 is false.
45004500
return (and_conditional_maps(left_if_vars, right_if_vars),
45014501
or_conditional_maps(left_else_vars, right_else_vars))
45024502
elif isinstance(node, OpExpr) and node.op == 'or':
4503-
left_if_vars, left_else_vars = self.find_isinstance_check_helper(node.left)
4504-
right_if_vars, right_else_vars = self.find_isinstance_check_helper(node.right)
4503+
left_if_vars, left_else_vars = self.find_isinstance_check(node.left)
4504+
right_if_vars, right_else_vars = self.find_isinstance_check(node.right)
45054505

45064506
# (e1 or e2) is true if at least one of e1 or e2 is true,
45074507
# and false if both e1 and e2 are false.
45084508
return (or_conditional_maps(left_if_vars, right_if_vars),
45094509
and_conditional_maps(left_else_vars, right_else_vars))
45104510
elif isinstance(node, UnaryExpr) and node.op == 'not':
4511-
left, right = self.find_isinstance_check_helper(node.expr)
4511+
left, right = self.find_isinstance_check(node.expr)
45124512
return right, left
45134513

45144514
# Restrict the type of the variable to True-ish/False-ish in the if and else branches

test-data/unit/check-narrowing.test

Lines changed: 20 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -655,6 +655,26 @@ else:
655655
reveal_type(y["model"]) # N: Revealed type is "Union[TypedDict('__main__.Model1', {'key': Literal['A']}), TypedDict('__main__.Model2', {'key': Literal['B']})]"
656656
[builtins fixtures/primitives.pyi]
657657

658+
[case testNarrowingExprPropagation]
659+
from typing import Union
660+
from typing_extensions import Literal
661+
662+
class A:
663+
tag: Literal['A']
664+
665+
class B:
666+
tag: Literal['B']
667+
668+
abo: Union[A, B, None]
669+
if abo is not None and abo.tag == "A":
670+
reveal_type(abo.tag) # N: Revealed type is "Literal['A']"
671+
reveal_type(abo) # N: Revealed type is "__main__.A"
672+
673+
if not (x := abo is None or abo.tag != "B"):
674+
reveal_type(abo.tag) # N: Revealed type is "Literal['B']"
675+
reveal_type(abo) # N: Revealed type is "__main__.B"
676+
[builtins fixtures/primitives.pyi]
677+
658678
[case testNarrowingEqualityFlipFlop]
659679
# flags: --warn-unreachable --strict-equality
660680
from typing_extensions import Literal, Final

0 commit comments

Comments
 (0)