Don't assume that for loop body index variable is always set (#18631)

Fixes #18629. Fixes #16321. Fixes #8637.

Author: JukkaL

mypy/checker.py

@@ -585,12 +585,12 @@ def accept_loop(
         else_body: Statement | None = None,
         *,
         exit_condition: Expression | None = None,
+        on_enter_body: Callable[[], None] | None = None,
     ) -> None:
         """Repeatedly type check a loop body until the frame doesn't change."""
 
         # The outer frame accumulates the results of all iterations:
         with self.binder.frame_context(can_skip=False, conditional_frame=True):
-
             # Check for potential decreases in the number of partial types so as not to stop the
             # iteration too early:
             partials_old = sum(len(pts.map) for pts in self.partial_types)
@@ -603,6 +603,9 @@ def accept_loop(
 
             while True:
                 with self.binder.frame_context(can_skip=True, break_frame=2, continue_frame=1):
+                    if on_enter_body is not None:
+                        on_enter_body()
+
                     self.accept(body)
                 partials_new = sum(len(pts.map) for pts in self.partial_types)
                 if (partials_new == partials_old) and not self.binder.last_pop_changed:
@@ -615,6 +618,9 @@ def accept_loop(
                 self.options.enabled_error_codes.add(codes.REDUNDANT_EXPR)
             if warn_unreachable or warn_redundant:
                 with self.binder.frame_context(can_skip=True, break_frame=2, continue_frame=1):
+                    if on_enter_body is not None:
+                        on_enter_body()
+
                     self.accept(body)
 
             # If exit_condition is set, assume it must be False on exit from the loop:
@@ -5126,8 +5132,14 @@ def visit_for_stmt(self, s: ForStmt) -> None:
             iterator_type, item_type = self.analyze_iterable_item_type(s.expr)
         s.inferred_item_type = item_type
         s.inferred_iterator_type = iterator_type
-        self.analyze_index_variables(s.index, item_type, s.index_type is None, s)
-        self.accept_loop(s.body, s.else_body)
+
+        self.accept_loop(
+            s.body,
+            s.else_body,
+            on_enter_body=lambda: self.analyze_index_variables(
+                s.index, item_type, s.index_type is None, s
+            ),
+        )
 
     def analyze_async_iterable_item_type(self, expr: Expression) -> tuple[Type, Type]:
         """Analyse async iterable expression and return iterator and iterator item types."""

test-data/unit/check-inference.test

@@ -3894,3 +3894,21 @@ foo = [
 ]
 reveal_type(foo)  # N: Revealed type is "builtins.list[Tuple[builtins.int, typing.Sequence[builtins.str]]]"
 [builtins fixtures/tuple.pyi]
+
+[case testForLoopIndexVaribaleNarrowing1]
+# flags: --local-partial-types
+from typing import Union
+x: Union[int, str]
+x = "abc"
+for x in list[int]():
+    reveal_type(x)  # N: Revealed type is "builtins.int"
+reveal_type(x) # N: Revealed type is "Union[builtins.int, builtins.str]"
+
+[case testForLoopIndexVaribaleNarrowing2]
+# flags: --enable-error-code=redundant-expr
+from typing import Union
+x: Union[int, str]
+x = "abc"
+for x in list[int]():
+    reveal_type(x)  # N: Revealed type is "builtins.int"
+reveal_type(x) # N: Revealed type is "Union[builtins.int, builtins.str]"