Switch error code used to report vars defined in different branch

mypy/partially_defined.py

@@ -234,6 +234,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.tracker = DefinedVariableTracker()
 
     def process_lvalue(self, lvalue: Lvalue | None) -> None:
@@ -319,10 +320,12 @@ 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
         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()
 
     def visit_return_stmt(self, o: ReturnStmt) -> None:
@@ -354,7 +357,9 @@ 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
         o.body.accept(self)
+        self.loop_depth -= 1
         if not checker.is_true_literal(o.expr):
             self.tracker.next_branch()
             if o.else_body:
@@ -380,7 +385,10 @@ 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.
-            self.msg.var_used_before_def(o.name, o)
+            if self.loop_depth > 0:
+                self.msg.variable_may_be_undefined(o.name, o)
+            else:
+                self.msg.var_used_before_def(o.name, o)
         elif self.tracker.is_undefined(o.name):
             # A variable is undefined. It could be due to two things:
             # 1. A variable is just totally undefined

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

@@ -206,6 +206,48 @@ def f5() -> int:
         return 3
     return 1
 
+[case testDefinedDifferentBranchUseBeforeDef]
+# flags: --enable-error-code partially-defined --enable-error-code use-before-def
+
+def f0() -> None:
+    if int():
+        x = 0
+    else:
+        y = x  # E: Name "x" is used before definition
+        z = x  # E: Name "x" is used before definition
+
+def f1() -> None:
+    x = 1
+    if int():
+        x = 0
+    else:
+        y = x  # No error.
+
+
+[case testDefinedDifferentBranchPartiallyDefined]
+# flags: --enable-error-code partially-defined --enable-error-code use-before-def
+
+def f0() -> None:
+    first_iter = True
+    for i in [0, 1]:
+        if first_iter:
+            first_iter = False
+            x = 0
+        else:
+            # This is technically a false positive but mypy isn't smart enough for this yet.
+            y = x  # E: Name "x" may be undefined
+            z = x  # E: Name "x" may be undefined
+
+
+def f1() -> None:
+    while True:
+        if int():
+            x = 0
+        else:
+            y = x  # E: Name "x" may be undefined
+            z = x  # E: Name "x" may be undefined
+
+
 [case testAssert]
 # flags: --enable-error-code partially-defined
 def f1() -> int:
@@ -394,21 +436,7 @@ def f0() -> None:
     x = y  # E: Name "y" is used before definition
     y: int = 1
 
-def f1() -> None:
-    if int():
-        x = 0
-    else:
-        y = x  # E: Name "x" is used before definition
-        z = x  # E: Name "x" is used before definition
-
 def f2() -> None:
-    x = 1
-    if int():
-        x = 0
-    else:
-        y = x  # No error.
-
-def f3() -> None:
     if int():
         pass
     else:
@@ -418,14 +446,14 @@ def f3() -> None:
     def inner2() -> None:
         z = 0
 
-def f4() -> None:
+def f3() -> None:
     if int():
         pass
     else:
         y = z  # E: Name "z" is used before definition
     z: int = 2
 
-def f5() -> None:
+def f4() -> None:
     if int():
         pass
     else: