[possibly-undefined] Fix loop misconceptions

mypy/partially_defined.py

@@ -396,10 +396,10 @@ def visit_if_stmt(self, o: IfStmt) -> None:
             b.accept(self)
             self.tracker.next_branch()
         if o.else_body:
-            if not o.else_body.is_unreachable:
-                o.else_body.accept(self)
-            else:
+            if o.else_body.is_unreachable:
                 self.tracker.skip_branch()
+            else:
+                o.else_body.accept(self)
         self.tracker.end_branch_statement()
 
     def visit_match_stmt(self, o: MatchStmt) -> None:
@@ -457,25 +457,27 @@ def visit_dictionary_comprehension(self, o: DictionaryComprehension) -> None:
 
     def visit_for_stmt(self, o: ForStmt) -> None:
         o.expr.accept(self)
+        self.tracker.start_branch_statement()
         self.process_lvalue(o.index)
         o.index.accept(self)
-        self.tracker.start_branch_statement()
         loop = Loop()
         self.loops.append(loop)
         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`.
-            has_break = loop.has_break
-            if has_break:
-                self.tracker.start_branch_statement()
-                self.tracker.next_branch()
+        if o.else_body is None:
+            self.tracker.end_branch_statement()
+        elif loop.has_break:
+            # `else` is executed conditionally:
+            # - if iterable was empty
+            # - if iterable was non-empty and `break` was not executed
+            o.else_body.accept(self)
+            self.tracker.end_branch_statement()
+        else:
+            # `else` is always executed:
+            # - if iterable was empty
+            # - if iterable was non-empty (there is no `break`)
+            self.tracker.end_branch_statement()
             o.else_body.accept(self)
-            if has_break:
-                self.tracker.end_branch_statement()
         self.loops.pop()
 
     def visit_return_stmt(self, o: ReturnStmt) -> None:
@@ -582,28 +584,34 @@ 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()
-            if o.else_body:
+
+            if o.else_body is None:
+                self.tracker.end_branch_statement()
+            elif loop.has_break:
+                # `else` is executed conditionally:
+                # - if expression was False on first evaluation
+                # - if expression was True on first evaluation and `break` was not executed
                 o.else_body.accept(self)
-            if has_break:
                 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

@@ -343,7 +343,12 @@ for x in [1, 2, 3]:
 else:
     z = 2
 
-a = z + y  # E: Name "y" may be undefined
+# mypy doesn't know if `[1, 2, 3]` is non-empty, so the body may not be executed
+x  # E: Name "x" may be undefined
+y  # E: Name "y" may be undefined
+
+# there is no `break`, so the `else` is always executed
+z
 
 [case testReturn]
 # flags: --enable-error-code possibly-undefined
@@ -1043,3 +1048,116 @@ def foo(x: Union[int, str]) -> None:
         assert_never(x)
     f  # OK
 [builtins fixtures/tuple.pyi]
+
+[case testForLoopIndexMayBeUndefined]
+# flags: --enable-error-code possibly-undefined
+class Iterable:
+    def __iter__(self): ...
+
+for i in Iterable():
+    pass
+
+i  # E: Name "i" may be undefined  # if iterable was empty
+
+[case testForLoopBodyMayNotBeExecuted]
+# flags: --enable-error-code possibly-undefined
+# Regression test for https://github.com/python/mypy/issues/14209
+class Iterable:
+    def __iter__(self): ...
+
+for _ in Iterable():
+    a = 1
+
+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
+class Iterable:
+    def __iter__(self): ...
+
+for _ in Iterable():
+    a = 1
+else:
+    a = 1  # if iterable is empty
+
+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
+class Iterable:
+    def __iter__(self): ...
+
+for i in Iterable():
+    a = 1
+    break
+else:
+    a = 1  # if iterable is empty
+
+a  # OK
+
+[case testForLoopBreakAsOnlyStatementOfCondition]
+# flags: --enable-error-code possibly-undefined
+# Regression test for https://github.com/python/mypy/issues/14209
+def condition() -> bool: ...
+class Iterable:
+    def __iter__(self): ...
+
+for i in Iterable():
+    a = 1
+    if condition():
+        break
+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