used context manager to manage flag, updated unit tests

Author: coldwolverine

mypy/semanal.py

@@ -485,8 +485,8 @@ def __init__(
         self.current_overload_item: int | None = None
 
         # Used to track whether currently inside an except* block. This helps
-        # to invoke errors when continue/break/return is used inside except * block.
-        self.is_in_except_block: bool = False
+        # to invoke errors when continue/break/return is used inside except* block.
+        self.inside_except_star_block: bool = False
 
     # mypyc doesn't properly handle implementing an abstractproperty
     # with a regular attribute so we make them properties
@@ -515,6 +515,15 @@ def allow_unbound_tvars_set(self) -> Iterator[None]:
         finally:
             self.allow_unbound_tvars = old
 
+    @contextmanager
+    def inside_except_star_block_set(self, value: bool) -> Iterator[None]:
+        old = self.inside_except_star_block
+        self.inside_except_star_block = value
+        try:
+            yield
+        finally:
+            self.inside_except_star_block = old
+
     #
     # Preparing module (performed before semantic analysis)
     #
@@ -858,9 +867,6 @@ def file_context(
     def visit_func_def(self, defn: FuncDef) -> None:
         self.statement = defn
 
-        was_in_except_block: bool = self.is_in_except_block
-        self.is_in_except_block = False
-
         # Visit default values because they may contain assignment expressions.
         for arg in defn.arguments:
             if arg.initializer:
@@ -881,13 +887,11 @@ def visit_func_def(self, defn: FuncDef) -> None:
                 self.add_function_to_symbol_table(defn)
 
         if not self.recurse_into_functions:
-            self.is_in_except_block = was_in_except_block
             return
 
         with self.scope.function_scope(defn):
-            self.analyze_func_def(defn)
-
-        self.is_in_except_block = was_in_except_block
+            with self.inside_except_star_block_set(False):
+                self.analyze_func_def(defn)
 
     def function_fullname(self, fullname: str) -> str:
         if self.current_overload_item is None:
@@ -5273,7 +5277,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.is_in_except_block:
+        if self.inside_except_star_block:
             self.fail('"return" not allowed in except* block', s, serious=True, blocker=True)
         if s.expr:
             s.expr.accept(self)
@@ -5307,12 +5311,10 @@ def visit_operator_assignment_stmt(self, s: OperatorAssignmentStmt) -> None:
     def visit_while_stmt(self, s: WhileStmt) -> None:
         self.statement = s
         s.expr.accept(self)
-        was_in_except_block: bool = self.is_in_except_block
-        self.is_in_except_block = False
         self.loop_depth[-1] += 1
-        s.body.accept(self)
+        with self.inside_except_star_block_set(False):
+            s.body.accept(self)
         self.loop_depth[-1] -= 1
-        self.is_in_except_block = was_in_except_block
         self.visit_block_maybe(s.else_body)
 
     def visit_for_stmt(self, s: ForStmt) -> None:
@@ -5334,27 +5336,24 @@ def visit_for_stmt(self, s: ForStmt) -> None:
                 self.store_declared_types(s.index, analyzed)
                 s.index_type = analyzed
 
-        was_in_except_block: bool = self.is_in_except_block
-        self.is_in_except_block = False
         self.loop_depth[-1] += 1
-        self.visit_block(s.body)
+        with self.inside_except_star_block_set(False):
+            self.visit_block(s.body)
         self.loop_depth[-1] -= 1
-        self.is_in_except_block = was_in_except_block
-
         self.visit_block_maybe(s.else_body)
 
     def visit_break_stmt(self, s: BreakStmt) -> None:
         self.statement = s
         if self.loop_depth[-1] == 0:
             self.fail('"break" outside loop', s, serious=True, blocker=True)
-        if self.is_in_except_block:
+        if self.inside_except_star_block:
             self.fail('"break" not allowed in except* block', s, serious=True, blocker=True)
 
     def visit_continue_stmt(self, s: ContinueStmt) -> None:
         self.statement = s
         if self.loop_depth[-1] == 0:
             self.fail('"continue" outside loop', s, serious=True, blocker=True)
-        if self.is_in_except_block:
+        if self.inside_except_star_block:
             self.fail('"continue" not allowed in except* block', s, serious=True, blocker=True)
 
     def visit_if_stmt(self, s: IfStmt) -> None:
@@ -5376,12 +5375,7 @@ def analyze_try_stmt(self, s: TryStmt, visitor: NodeVisitor[None]) -> None:
                 type.accept(visitor)
             if var:
                 self.analyze_lvalue(var)
-            if s.is_star:
-                was_in_except_block: bool = self.is_in_except_block
-                self.is_in_except_block = True
-                handler.accept(visitor)
-                self.is_in_except_block = was_in_except_block
-            else:
+            with self.inside_except_star_block_set(s.is_star):
                 handler.accept(visitor)
         if s.else_body:
             s.else_body.accept(visitor)

test-data/unit/check-except-star.test

@@ -173,3 +173,52 @@ Alias4 = Callable[[*IntList], int]  # E: "List[int]" cannot be unpacked (must be
 x4: Alias4[int]  # E: Bad number of arguments for type alias, expected 0, given 1
 reveal_type(x4)  # N: Revealed type is "def (*Any) -> builtins.int"
 [builtins fixtures/tuple.pyi]
+
+[case testReturnInExceptStarBlock]
+# flags: --python-version 3.11
+def foo() -> None:
+    try:
+        pass
+    except* Exception:
+        return   # E: "return" not allowed in except* block
+    finally:
+        return
+[builtins fixtures/exception.pyi]
+
+[case testBreakContinueReturnInExceptStarBlock1]
+# flags: --python-version 3.11
+class range: pass
+def foo() -> None:
+    for _ in range(5):
+        try:
+            pass
+        except* Exception:
+            continue   # E: "continue" not allowed in except* block
+        except* Exception:
+            for _ in range(2):
+                continue
+            break  # E: "break" not allowed in except* block
+        except* Exception:
+            return  # E: "return" not allowed in except* block
+[builtins fixtures/exception.pyi]
+
+[case testBreakContinueReturnInExceptStarBlock2]
+# flags: --python-version 3.11
+def foo():
+    while True:
+        try:
+            pass
+        except* Exception:
+            def inner():
+                while True:
+                    if 1 < 1:
+                        continue
+                    else:
+                        break
+                return
+            if 1 < 2:
+                break  # E: "break" not allowed in except* block
+            if 1 < 2:
+                continue  # E: "continue" not allowed in except* block
+            return  # E: "return" not allowed in except* block
+[builtins fixtures/exception.pyi]