add except* support to 102, 103, 104, 120, 910, 911 & 912 (#358)

* add except* support to 102, 103, 104, 120, 910, 911 & 912

Author: jakkdl

.pre-commit-config.yaml

@@ -42,7 +42,7 @@ repos:
       hooks:
           - id: mypy
             # uses py311 syntax, mypy configured for py39
-            exclude: tests/eval_files/.*_py311.py
+            exclude: tests/(eval|autofix)_files/.*_py311.py
 
     - repo: https://github.com/RobertCraigie/pyright-python
       rev: v1.1.396

docs/changelog.rst

@@ -4,6 +4,10 @@ Changelog
 
 `CalVer, YY.month.patch <https://calver.org/>`_
 
+25.3.1
+======
+- Add except* support to ASYNC102, 103, 104, 120, 910, 911, 912.
+
 25.2.3
 =======
 - No longer require ``flake8`` for installation... so if you require support for config files you must install ``flake8-async[flake8]``

docs/usage.rst

@@ -33,7 +33,7 @@ adding the following to your ``.pre-commit-config.yaml``:
    minimum_pre_commit_version: '2.9.0'
    repos:
    - repo: https://github.com/python-trio/flake8-async
-     rev: 25.2.3
+     rev: 25.3.1
      hooks:
        - id: flake8-async
          # args: ["--enable=ASYNC100,ASYNC112", "--disable=", "--autofix=ASYNC"]

flake8_async/__init__.py

@@ -38,7 +38,7 @@
 
 
 # CalVer: YY.month.patch, e.g. first release of July 2022 == "22.7.1"
-__version__ = "25.2.3"
+__version__ = "25.3.1"
 
 
 # taken from https://github.com/Zac-HD/shed

flake8_async/visitors/visitor102_120.py

@@ -149,7 +149,7 @@ def visit_AsyncWith(self, node: ast.AsyncWith):
                 break
         self.visit_With(node)
 
-    def visit_Try(self, node: ast.Try):
+    def visit_Try(self, node: ast.Try | ast.TryStar):  # type: ignore[name-defined]
         self.save_state(
             node, "_critical_scope", "_trio_context_managers", "cancelled_caught"
         )
@@ -165,6 +165,8 @@ def visit_Try(self, node: ast.Try):
         self._critical_scope = Statement("try/finally", node.lineno, node.col_offset)
         self.visit_nodes(node.finalbody)
 
+    visit_TryStar = visit_Try
+
     def visit_ExceptHandler(self, node: ast.ExceptHandler):
         # if we're inside a critical scope, a nested except should never override that
         if self._critical_scope is not None and self._critical_scope.name != "except":

flake8_async/visitors/visitor103_104.py

@@ -160,7 +160,7 @@ def visit_Return(self, node: ast.Return | ast.Yield):
     visit_Yield = visit_Return
 
     # Treat Try's as fully covering only if `finally` always raises.
-    def visit_Try(self, node: ast.Try):
+    def visit_Try(self, node: ast.Try | ast.TryStar):  # type: ignore[name-defined]
         self.save_state(node, "cancelled_caught", copy=True)
         self.cancelled_caught = set()
 
@@ -179,6 +179,8 @@ def visit_Try(self, node: ast.Try):
         # but it's fine if we raise in finally
         self.visit_nodes(node.finalbody)
 
+    visit_TryStar = visit_Try
+
     # Treat if's as fully covering if both `if` and `else` raise.
     # `elif` is parsed by the ast as a new if statement inside the else.
     def visit_If(self, node: ast.If):

flake8_async/visitors/visitor91x.py

@@ -768,7 +768,7 @@ def leave_Yield(
     # try can jump into any except or into the finally* at any point during it's
     # execution so we need to make sure except & finally can handle worst-case
     # * unless there's a bare except / except BaseException - not implemented.
-    def visit_Try(self, node: cst.Try):
+    def visit_Try(self, node: cst.Try | cst.TryStar):
         if not self.async_function:
             return
         self.save_state(node, "try_state", copy=True)
@@ -784,39 +784,41 @@ def visit_Try(self, node: cst.Try):
                 Statement("yield", pos.line, pos.column)  # type: ignore
             )
 
-    def leave_Try_body(self, node: cst.Try):
+    def leave_Try_body(self, node: cst.Try | cst.TryStar):
         # save state at end of try for entering else
         self.try_state.try_checkpoint = self.uncheckpointed_statements
 
         # check that all except handlers checkpoint (await or most likely raise)
         self.try_state.except_uncheckpointed_statements = set()
 
-    def visit_ExceptHandler(self, node: cst.ExceptHandler):
+    def visit_ExceptHandler(self, node: cst.ExceptHandler | cst.ExceptStarHandler):
         # enter with worst case of try
         self.uncheckpointed_statements = (
             self.try_state.body_uncheckpointed_statements.copy()
         )
 
     def leave_ExceptHandler(
-        self, original_node: cst.ExceptHandler, updated_node: cst.ExceptHandler
-    ) -> cst.ExceptHandler:
+        self,
+        original_node: cst.ExceptHandler | cst.ExceptStarHandler,
+        updated_node: cst.ExceptHandler | cst.ExceptStarHandler,
+    ) -> Any:  # not worth creating a TypeVar to handle correctly
         self.try_state.except_uncheckpointed_statements.update(
             self.uncheckpointed_statements
         )
         return updated_node
 
-    def visit_Try_orelse(self, node: cst.Try):
+    def visit_Try_orelse(self, node: cst.Try | cst.TryStar):
         # check else
         # if else runs it's after all of try, so restore state to back then
         self.uncheckpointed_statements = self.try_state.try_checkpoint
 
-    def leave_Try_orelse(self, node: cst.Try):
+    def leave_Try_orelse(self, node: cst.Try | cst.TryStar):
         # checkpoint if else checkpoints, and all excepts checkpoint
         self.uncheckpointed_statements.update(
             self.try_state.except_uncheckpointed_statements
         )
 
-    def visit_Try_finalbody(self, node: cst.Try):
+    def visit_Try_finalbody(self, node: cst.Try | cst.TryStar):
         if node.finalbody:
             self.try_state.added = (
                 self.try_state.body_uncheckpointed_statements.difference(
@@ -835,14 +837,26 @@ def visit_Try_finalbody(self, node: cst.Try):
             ):
                 self.uncheckpointed_statements.update(self.try_state.added)
 
-    def leave_Try_finalbody(self, node: cst.Try):
+    def leave_Try_finalbody(self, node: cst.Try | cst.TryStar):
         if node.finalbody:
             self.uncheckpointed_statements.difference_update(self.try_state.added)
 
-    def leave_Try(self, original_node: cst.Try, updated_node: cst.Try) -> cst.Try:
+    def leave_Try(
+        self, original_node: cst.Try | cst.TryStar, updated_node: cst.Try | cst.TryStar
+    ) -> cst.Try | cst.TryStar:
         self.restore_state(original_node)
         return updated_node
 
+    visit_TryStar = visit_Try
+    leave_TryStar = leave_Try
+    leave_TryStar_body = leave_Try_body
+    visit_TryStar_orelse = visit_Try_orelse
+    leave_TryStar_orelse = leave_Try_orelse
+    visit_TryStar_finalbody = visit_Try_finalbody
+    leave_TryStar_finalbody = leave_Try_finalbody
+    visit_ExceptStarHandler = visit_ExceptHandler
+    leave_ExceptStarHandler = leave_ExceptHandler
+
     def leave_If_test(self, node: cst.If | cst.IfExp) -> None:
         if not self.async_function:
             return

tests/autofix_files/async91x_py311.py

@@ -0,0 +1,97 @@
+"""Test for ASYNC91x rules with except* blocks.
+
+ASYNC910: async-function-no-checkpoint
+ASYNC911: async-generator-no-checkpoint
+ASYNC913: indefinite-loop-no-guaranteed-checkpoint
+
+async912 handled in separate file
+"""
+
+# ARG --enable=ASYNC910,ASYNC911,ASYNC913
+# AUTOFIX
+# ASYNCIO_NO_AUTOFIX
+import trio
+
+
+async def foo(): ...
+
+
+async def foo_try_except_star_1():  # ASYNC910: 0, "exit", Statement("function definition", lineno)
+    try:
+        await foo()
+    except* ValueError:
+        ...
+    except* RuntimeError:
+        raise
+    else:
+        await foo()
+    await trio.lowlevel.checkpoint()
+
+
+async def foo_try_except_star_2():  # safe
+    try:
+        ...
+    except* ValueError:
+        ...
+    finally:
+        await foo()
+
+
+async def foo_try_except_star_3():  # safe
+    try:
+        await foo()
+    except* ValueError:
+        raise
+
+
+# Multiple except* handlers - should all guarantee checkpoint/raise
+async def foo_try_except_star_4():
+    try:
+        await foo()
+    except* ValueError:
+        await foo()
+    except* TypeError:
+        raise
+    except* Exception:
+        raise
+
+
+async def try_else_no_raise_in_except():  # ASYNC910: 0, "exit", Statement("function definition", lineno)
+    try:
+        ...
+    except* ValueError:
+        ...
+    else:
+        await foo()
+    await trio.lowlevel.checkpoint()
+
+
+async def try_else_raise_in_except():
+    try:
+        ...
+    except* ValueError:
+        raise
+    else:
+        await foo()
+
+
+async def check_async911():  # ASYNC911: 0, "exit", Statement("yield", lineno+7)
+    try:
+        await foo()
+    except* ValueError:
+        ...
+    except* RuntimeError:
+        raise
+    await trio.lowlevel.checkpoint()
+    yield  # ASYNC911: 4, "yield", Statement("function definition", lineno-7)
+    await trio.lowlevel.checkpoint()
+
+
+async def check_async913():
+    while True:  # ASYNC913: 4
+        await trio.lowlevel.checkpoint()
+        try:
+            await foo()
+        except* ValueError:
+            # Missing checkpoint
+            ...

tests/autofix_files/async91x_py311.py.diff

@@ -0,0 +1,33 @@
+---
++++
+@@ x,6 x,7 @@
+         raise
+     else:
+         await foo()
++    await trio.lowlevel.checkpoint()
+
+
+ async def foo_try_except_star_2():  # safe
+@@ x,6 x,7 @@
+         ...
+     else:
+         await foo()
++    await trio.lowlevel.checkpoint()
+
+
+ async def try_else_raise_in_except():
+@@ x,11 x,14 @@
+         ...
+     except* RuntimeError:
+         raise
++    await trio.lowlevel.checkpoint()
+     yield  # ASYNC911: 4, "yield", Statement("function definition", lineno-7)
++    await trio.lowlevel.checkpoint()
+
+
+ async def check_async913():
+     while True:  # ASYNC913: 4
++        await trio.lowlevel.checkpoint()
+         try:
+             await foo()
+         except* ValueError:

tests/eval_files/async102_120_py311.py

@@ -0,0 +1,29 @@
+"""Test for ASYNC102/ASYNC120 with except*
+
+ASYNC102: await-in-finally-or-cancelled
+
+ASYNC120: await-in-except
+"""
+
+# type: ignore
+# ARG --enable=ASYNC102,ASYNC120
+# NOASYNCIO # TODO: support asyncio shields
+import trio
+
+
+async def foo():
+    try:
+        ...
+    except* ValueError:
+        await foo()  # ASYNC120: 8, Statement("except", lineno-1)
+        raise
+    except* BaseException:
+        await foo()  # ASYNC102: 8, Statement("BaseException", lineno-1)
+    finally:
+        await foo()  # ASYNC102: 8, Statement("try/finally", lineno-8)
+
+    try:
+        ...
+    except* BaseException:
+        with trio.move_on_after(30, shield=True):
+            await foo()