diff --git a/mypy/checker.py b/mypy/checker.py index 999d75678aa4..09bfb4e0f568 100644 --- a/mypy/checker.py +++ b/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: @@ -5119,8 +5125,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.""" diff --git a/test-data/unit/check-inference.test b/test-data/unit/check-inference.test index 473a3f9d3df6..d80181047dc8 100644 --- a/test-data/unit/check-inference.test +++ b/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]"