added check for edge case when return stmt inside loop can affect control flow in except* block

Author: coldwolverine

mypy/semanal.py

@@ -487,6 +487,8 @@ def __init__(
         # Used to track whether currently inside an except* block. This helps
         # to invoke errors when continue/break/return is used inside except* block.
         self.inside_except_star_block: bool = False
+        # Used to track edge case when return is still inside except* if it enters a loop
+        self.return_stmt_inside_except_star_block: bool = False
 
     # mypyc doesn't properly handle implementing an abstractproperty
     # with a regular attribute so we make them properties
@@ -516,13 +518,23 @@ def allow_unbound_tvars_set(self) -> Iterator[None]:
             self.allow_unbound_tvars = old
 
     @contextmanager
-    def inside_except_star_block_set(self, value: bool) -> Iterator[None]:
+    def inside_except_star_block_set(
+        self, value: bool, enteringLoop: bool = False
+    ) -> Iterator[None]:
         old = self.inside_except_star_block
         self.inside_except_star_block = value
+
+        # Return statement would still be in except* scope if entering loops
+        if not enteringLoop:
+            old_return_stmt_flag = self.return_stmt_inside_except_star_block
+            self.return_stmt_inside_except_star_block = value
+
         try:
             yield
         finally:
             self.inside_except_star_block = old
+            if not enteringLoop:
+                self.return_stmt_inside_except_star_block = old_return_stmt_flag
 
     #
     # Preparing module (performed before semantic analysis)
@@ -890,7 +902,7 @@ def visit_func_def(self, defn: FuncDef) -> None:
             return
 
         with self.scope.function_scope(defn):
-            with self.inside_except_star_block_set(False):
+            with self.inside_except_star_block_set(value=False):
                 self.analyze_func_def(defn)
 
     def function_fullname(self, fullname: str) -> str:
@@ -5277,7 +5289,7 @@ def visit_return_stmt(self, s: ReturnStmt) -> None:
         self.statement = s
         if not self.is_func_scope():
             self.fail('"return" outside function', s)
-        if self.inside_except_star_block:
+        if self.return_stmt_inside_except_star_block:
             self.fail('"return" not allowed in except* block', s, serious=True, blocker=True)
         if s.expr:
             s.expr.accept(self)
@@ -5312,7 +5324,7 @@ def visit_while_stmt(self, s: WhileStmt) -> None:
         self.statement = s
         s.expr.accept(self)
         self.loop_depth[-1] += 1
-        with self.inside_except_star_block_set(False):
+        with self.inside_except_star_block_set(value=False, enteringLoop=True):
             s.body.accept(self)
         self.loop_depth[-1] -= 1
         self.visit_block_maybe(s.else_body)
@@ -5337,7 +5349,7 @@ def visit_for_stmt(self, s: ForStmt) -> None:
                 s.index_type = analyzed
 
         self.loop_depth[-1] += 1
-        with self.inside_except_star_block_set(False):
+        with self.inside_except_star_block_set(value=False, enteringLoop=True):
             self.visit_block(s.body)
         self.loop_depth[-1] -= 1
         self.visit_block_maybe(s.else_body)

test-data/unit/check-python311.test

@@ -174,7 +174,7 @@ x4: Alias4[int]  # E: Bad number of arguments for type alias, expected 0, given
 reveal_type(x4)  # N: Revealed type is "def (*Any) -> builtins.int"
 [builtins fixtures/tuple.pyi]
 
-[case testReturnInExceptStarBlock]
+[case testReturnInExceptStarBlock1]
 # flags: --python-version 3.11
 def foo() -> None:
     try:
@@ -185,6 +185,17 @@ def foo() -> None:
         return
 [builtins fixtures/exception.pyi]
 
+[case testReturnInExceptStarBlock2]
+# flags: --python-version 3.11
+def foo():
+    while True:
+        try:
+            pass
+        except* Exception:
+            while True:
+                return  # E: "return" not allowed in except* block
+[builtins fixtures/exception.pyi]
+
 [case testBreakContinueReturnInExceptStarBlock1]
 # flags: --python-version 3.11
 class range: pass