[partially defined] rework handling of else in loops (#14191)

`else` in loops is executed if the loop didn't exit via a raise, return,
or a break. Therefore, we should treat it execution as unconditional if
there is not a break statement.

This PR also improves error messages in `else`, reporting them as "may
be undefined" instead of "used before definition"

Fixes #14097

Co-authored-by: Ivan Levkivskyi <[email protected]>

mypy/partially_defined.py

@@ -224,6 +224,11 @@ def refers_to_builtin(o: RefExpr) -> bool:
     return o.fullname is not None and o.fullname.startswith("builtins.")
 
 
+class Loop:
+    def __init__(self) -> None:
+        self.has_break = False
+
+
 class PartiallyDefinedVariableVisitor(ExtendedTraverserVisitor):
     """Detects the following cases:
     - A variable that's defined only part of the time.
@@ -245,7 +250,7 @@ class PartiallyDefinedVariableVisitor(ExtendedTraverserVisitor):
     def __init__(self, msg: MessageBuilder, type_map: dict[Expression, Type]) -> None:
         self.msg = msg
         self.type_map = type_map
-        self.loop_depth = 0
+        self.loops: list[Loop] = []
         self.tracker = DefinedVariableTracker()
         for name in implicit_module_attrs:
             self.tracker.record_definition(name)
@@ -344,13 +349,23 @@ def visit_for_stmt(self, o: ForStmt) -> None:
         self.process_lvalue(o.index)
         o.index.accept(self)
         self.tracker.start_branch_statement()
-        self.loop_depth += 1
+        loop = Loop()
+        self.loops.append(loop)
         o.body.accept(self)
         self.tracker.next_branch()
-        if o.else_body:
-            o.else_body.accept(self)
-        self.loop_depth -= 1
         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()
+            o.else_body.accept(self)
+            if has_break:
+                self.tracker.end_branch_statement()
+        self.loops.pop()
 
     def visit_return_stmt(self, o: ReturnStmt) -> None:
         super().visit_return_stmt(o)
@@ -376,6 +391,8 @@ def visit_continue_stmt(self, o: ContinueStmt) -> None:
 
     def visit_break_stmt(self, o: BreakStmt) -> None:
         super().visit_break_stmt(o)
+        if self.loops:
+            self.loops[-1].has_break = True
         self.tracker.skip_branch()
 
     def visit_expression_stmt(self, o: ExpressionStmt) -> None:
@@ -386,14 +403,28 @@ def visit_expression_stmt(self, o: ExpressionStmt) -> None:
     def visit_while_stmt(self, o: WhileStmt) -> None:
         o.expr.accept(self)
         self.tracker.start_branch_statement()
-        self.loop_depth += 1
+        loop = Loop()
+        self.loops.append(loop)
         o.body.accept(self)
-        self.loop_depth -= 1
+        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.
             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:
                 o.else_body.accept(self)
-        self.tracker.end_branch_statement()
+            if has_break:
+                self.tracker.end_branch_statement()
+        self.loops.pop()
 
     def visit_as_pattern(self, o: AsPattern) -> None:
         if o.name is not None:
@@ -415,7 +446,7 @@ def visit_name_expr(self, o: NameExpr) -> None:
             self.tracker.record_definition(o.name)
         elif self.tracker.is_defined_in_different_branch(o.name):
             # A variable is defined in one branch but used in a different branch.
-            if self.loop_depth > 0:
+            if self.loops:
                 self.variable_may_be_undefined(o.name, o)
             else:
                 self.var_used_before_def(o.name, o)

test-data/unit/check-partially-defined.test

@@ -167,30 +167,56 @@ def foo(a: int) -> None:
 [case testWhile]
 # flags: --enable-error-code partially-defined
 while int():
-    x = 1
+    a = 1
 
-y = x  # E: Name "x" may be undefined
+x = a  # E: Name "a" may be undefined
 
 while int():
-    z = 1
+    b = 1
 else:
-    z = 2
+    b = 2
 
-y = z  # No error.
+y = b  # No error.
 
 while True:
-    k = 1
+    c = 1
     if int():
         break
-y = k # No error.
+y = c  # No error.
+
+# This while loop doesn't have a `break` inside, so we know that the else must always get executed.
+while int():
+    pass
+else:
+    d = 1
+y = d  # No error.
+
+while int():
+    if int():
+        break
+else:
+    e = 1
+# If a while loop has a `break`, it's possible that the else didn't get executed.
+y = e  # E: Name "e" may be undefined
+
+while int():
+    while int():
+        if int():
+            break
+    else:
+        f = 1
+else:
+    g = 2
+
+y = f  # E: Name "f" may be undefined
+y = g
 
 [case testForLoop]
 # flags: --enable-error-code partially-defined
 for x in [1, 2, 3]:
     if x:
         x = 1
     y = x
-    z = 1
 else:
     z = 2
 
@@ -283,6 +309,17 @@ def f1() -> None:
             y = x  # E: Name "x" may be undefined
             z = x  # E: Name "x" may be undefined
 
+def f2() -> None:
+    for i in [0, 1]:
+        x = i
+    else:
+        y = x  # E: Name "x" may be undefined
+
+def f3() -> None:
+    while int():
+        x = 1
+    else:
+        y = x  # E: Name "x" may be undefined
 
 [case testAssert]
 # flags: --enable-error-code partially-defined
@@ -338,16 +375,12 @@ def f2() -> int:
     while int():
         if int():
             x = 1
-            z = 1
         elif int():
             pass
         else:
             continue
         y = x  # E: Name "x" may be undefined
-    else:
-        x = 2
-        z = 2
-    return z  # E: Name "z" may be undefined
+    return x  # E: Name "x" may be undefined
 
 def f3() -> None:
     while True: