Skip to content

Commit e8c39c9

Browse files
committed
Learn that loop condition is False on exit from while loop
1 parent 29b9f46 commit e8c39c9

File tree

2 files changed

+48
-2
lines changed

2 files changed

+48
-2
lines changed

mypy/checker.py

Lines changed: 12 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -226,8 +226,10 @@ def accept(self, node: Node, type_context: Type = None) -> Type:
226226
else:
227227
return typ
228228

229-
def accept_loop(self, body: Node, else_body: Node = None) -> Type:
229+
def accept_loop(self, body: Node, else_body: Node = None, *,
230+
exit_condition: Node = None) -> Type:
230231
"""Repeatedly type check a loop body until the frame doesn't change.
232+
If exit_condition is set, assume it must be False on exit from the loop.
231233
232234
Then check the else_body.
233235
"""
@@ -240,6 +242,13 @@ def accept_loop(self, body: Node, else_body: Node = None) -> Type:
240242
if not self.binder.last_pop_changed:
241243
break
242244
self.binder.pop_loop_frame()
245+
if exit_condition:
246+
_, else_map = find_isinstance_check(
247+
exit_condition, self.type_map, self.typing_mode_weak()
248+
)
249+
if else_map:
250+
for var, type in else_map.items():
251+
self.binder.push(var, type)
243252
if else_body:
244253
self.accept(else_body)
245254

@@ -1465,7 +1474,8 @@ def visit_if_stmt(self, s: IfStmt) -> Type:
14651474

14661475
def visit_while_stmt(self, s: WhileStmt) -> Type:
14671476
"""Type check a while statement."""
1468-
self.accept_loop(IfStmt([s.expr], [s.body], None), s.else_body)
1477+
self.accept_loop(IfStmt([s.expr], [s.body], None), s.else_body,
1478+
exit_condition=s.expr)
14691479

14701480
def visit_operator_assignment_stmt(self,
14711481
s: OperatorAssignmentStmt) -> Type:

test-data/unit/check-isinstance.test

Lines changed: 36 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -908,6 +908,42 @@ def bar() -> None:
908908
[out]
909909
main: note: In function "bar":
910910

911+
[case testWhileExitCondition1]
912+
from typing import Union
913+
x = 1 # type: Union[int, str]
914+
while isinstance(x, int):
915+
if bool():
916+
continue
917+
x = 'a'
918+
else:
919+
reveal_type(x) # E: Revealed type is 'builtins.str'
920+
reveal_type(x) # E: Revealed type is 'builtins.str'
921+
[builtins fixtures/isinstance.py]
922+
923+
[case testWhileExitCondition2]
924+
from typing import Union
925+
x = 1 # type: Union[int, str]
926+
while isinstance(x, int):
927+
if bool():
928+
break
929+
x = 'a'
930+
else:
931+
reveal_type(x) # E: Revealed type is 'builtins.str'
932+
reveal_type(x) # E: Revealed type is 'Union[builtins.int, builtins.str]'
933+
[builtins fixtures/isinstance.py]
934+
935+
[case testWhileLinkedList]
936+
from typing import Union
937+
LinkedList = Union['Cons', 'Nil']
938+
class Nil: pass
939+
class Cons:
940+
tail = None # type: LinkedList
941+
def last(x: LinkedList) -> Nil:
942+
while isinstance(x, Cons):
943+
x = x.tail
944+
return x
945+
[builtins fixtures/isinstance.py]
946+
911947
[case testReturnAndFlow]
912948
def foo() -> int:
913949
return 1 and 2

0 commit comments

Comments
 (0)