Fix: while loop else body is executed if first condition is False

mypy/partially_defined.py

@@ -576,27 +576,35 @@ def process_try_stmt(self, o: TryStmt) -> None:
 
     def visit_while_stmt(self, o: WhileStmt) -> None:
         o.expr.accept(self)
-        self.tracker.start_branch_statement()
         loop = Loop()
         self.loops.append(loop)
-        o.body.accept(self)
-        has_break = loop.has_break
-        if not checker.is_true_literal(o.expr):
-            # If this is a loop like `while True`, we can consider the body to be
-            # a single branch statement (we're guaranteed that the body is executed at least once).
-            # If not, call next_branch() to make all variables defined there conditional.
+
+        if checker.is_true_literal(o.expr):
+            # `while True` loop:
+            # - body is always executed
+            # - `else` is never executed
+            o.body.accept(self)
+        else:
+            # body is executed conditionally: if expression was True on first evaluation
+            self.tracker.start_branch_statement()
+            o.body.accept(self)
             self.tracker.next_branch()
-        self.tracker.end_branch_statement()
-        if o.else_body is not None:
-            # If the loop has a `break` inside, `else` is executed conditionally.
-            # If the loop doesn't have a `break` either the function will return or
-            # execute the `else`.
-            if has_break:
-                self.tracker.start_branch_statement()
-                self.tracker.next_branch()
-            o.else_body.accept(self)
-            if has_break:
+
+            if o.else_body is None:
                 self.tracker.end_branch_statement()
+            else:
+                if loop.has_break:
+                    # `else` is executed conditionally:
+                    # - if expression was False on first evaluation
+                    # - if `break` was not executed
+                    o.else_body.accept(self)
+                    self.tracker.end_branch_statement()
+                else:
+                    # `else` is always executed:
+                    # - if expression was False on first evaluation
+                    # - if expression was True on first evaluation (there is no break)
+                    self.tracker.end_branch_statement()
+                    o.else_body.accept(self)
         self.loops.pop()
 
     def visit_as_pattern(self, o: AsPattern) -> None:

test-data/unit/check-possibly-undefined.test

@@ -1069,6 +1069,16 @@ for _ in Iterable():
 
 a  # E: Name "a" may be undefined  # if iterable was empty
 
+[case testWhileLoopBodyMayNotBeExecuted]
+# flags: --enable-error-code possibly-undefined
+# Regression test for https://github.com/python/mypy/issues/14209
+def condition() -> bool: ...
+
+while condition():
+    a = 1
+
+a  # E: Name "a" may be undefined  # if `condition()` was `False` on first call
+
 [case testForLoopBodyOrElseIsExecuted]
 # flags: --enable-error-code possibly-undefined
 # Regression test for https://github.com/python/mypy/issues/14209
@@ -1082,6 +1092,18 @@ else:
 
 a  # OK
 
+[case testWhileLoopBodyOrElseIsExecuted]
+# flags: --enable-error-code possibly-undefined
+# Regression test for https://github.com/python/mypy/issues/14209
+def condition() -> bool: ...
+
+while condition():
+    a = 1
+else:
+    a = 1  # if `condition()` was `False` on first call
+
+a  # OK
+
 [case testForLoopBreakAsLastStatement]
 # flags: --enable-error-code possibly-undefined
 # Regression test for https://github.com/python/mypy/issues/14209
@@ -1111,3 +1133,30 @@ else:
     a = 1  # if iterable is empty
 
 a  # OK
+
+[case testWhileLoopBreakAsLastStatement]
+# flags: --enable-error-code possibly-undefined
+# Regression test for https://github.com/python/mypy/issues/14209
+def condition() -> bool: ...
+
+while condition():
+    a = 1
+    break
+else:
+    a = 1  # if `condition()` was `False` on first call
+
+a  # OK
+
+[case testWhileLoopBreakAsOnlyStatementOfCondition]
+# flags: --enable-error-code possibly-undefined
+# Regression test for https://github.com/python/mypy/issues/14209
+def condition() -> bool: ...
+
+while condition():
+    a = 1
+    if condition():
+        break
+else:
+    a = 1  # if `condition()` was `False` on first call
+
+a  # OK