add async120, await-in-except (#265)

* add async120, await-in-except
* Fix false alarm on nested function definitions for async102
---------

Co-authored-by: Zac Hatfield-Dodds <[email protected]>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

docs/changelog.rst

@@ -4,6 +4,11 @@ Changelog
 
 *[CalVer, YY.month.patch](https://calver.org/)*
 
+24.6.1
+======
+- Add :ref:`ASYNC120 <async120>` await-in-except.
+- Fix false alarm with :ref:`ASYNC102 <async102>` with function definitions inside finally/except.
+
 24.5.6
 ======
 - Make :ref:`ASYNC913 <async913>` disabled by default, as originally intended.

docs/rules.rst

@@ -20,8 +20,10 @@ ASYNC101 : yield-in-cancel-scope
     See `this thread <https://discuss.python.org/t/preventing-yield-inside-certain-context-managers/1091/23>`_ for discussion of a future PEP.
     This has substantial overlap with :ref:`ASYNC119 <ASYNC119>`, which will warn on almost all instances of ASYNC101, but ASYNC101 is about a conceptually different problem that will not get resolved by `PEP 533 <https://peps.python.org/pep-0533/>`_.
 
-ASYNC102 : await-in-finally-or-cancelled
+_`ASYNC102` : await-in-finally-or-cancelled
     ``await`` inside ``finally`` or :ref:`cancelled-catching <cancelled>` ``except:`` must have shielded :ref:`cancel scope <cancel_scope>` with timeout.
+    If not, the async call will immediately raise a new cancellation, suppressing the cancellation that was caught.
+    See :ref:`ASYNC120 <async120>` for the general case where other exceptions might get suppressed.
     This is currently not able to detect asyncio shields.
 
 ASYNC103 : no-reraise-cancelled
@@ -75,6 +77,12 @@ _`ASYNC119` : yield-in-cm-in-async-gen
    ``yield`` in context manager in async generator is unsafe, the cleanup may be delayed until ``await`` is no longer allowed.
    We strongly encourage you to read `PEP 533 <https://peps.python.org/pep-0533/>`_ and use `async with aclosing(...) <https://docs.python.org/3/library/contextlib.html#contextlib.aclosing>`_, or better yet avoid async generators entirely (see `ASYNC900`_ ) in favor of context managers which return an iterable :ref:`channel/stream/queue <channel_stream_queue>`.
 
+_`ASYNC120` : await-in-except
+    Dangerous :ref:`checkpoint` inside an ``except`` block.
+    If this checkpoint is cancelled, the current active exception will be replaced by the ``Cancelled`` exception, and cannot be reraised later.
+    This will not trigger when :ref:`ASYNC102 <ASYNC102>` does, and if you don't care about losing non-cancelled exceptions you could disable this rule.
+    This is currently not able to detect asyncio shields.
+
 
 Blocking sync calls in async functions
 ======================================

flake8_async/__init__.py

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

flake8_async/visitors/visitor102.py

@@ -23,6 +23,10 @@ class Visitor102(Flake8AsyncVisitor):
             "await inside {0.name} on line {0.lineno} must have shielded cancel "
             "scope with a timeout."
         ),
+        "ASYNC120": (
+            "checkpoint inside {0.name} on line {0.lineno} will discard the active "
+            "exception if cancelled."
+        ),
     }
 
     class TrioScope:
@@ -52,6 +56,15 @@ def __init__(self, *args: Any, **kwargs: Any):
         self._trio_context_managers: list[Visitor102.TrioScope] = []
         self.cancelled_caught = False
 
+        # list of dangerous awaits inside a non-critical except handler,
+        # which will emit errors upon reaching a raise.
+        # Entries are only added to the list inside except handlers,
+        # so with the `save_state` in visit_ExceptHandler any entries not followed
+        # by a raise will be thrown out when exiting the except handler.
+        self._potential_120: list[
+            tuple[ast.Await | ast.AsyncFor | ast.AsyncWith, Statement]
+        ] = []
+
     # if we're inside a finally or critical except, and we're not inside a scope that
     # doesn't have both a timeout and shield
     def async_call_checker(
@@ -60,7 +73,16 @@ def async_call_checker(
         if self._critical_scope is not None and not any(
             cm.has_timeout and cm.shielded for cm in self._trio_context_managers
         ):
-            self.error(node, self._critical_scope)
+            # non-critical exception handlers have the statement name set to "except"
+            if self._critical_scope.name == "except":
+                self._potential_120.append((node, self._critical_scope))
+            else:
+                self.error(node, self._critical_scope, error_code="ASYNC102")
+
+    def visit_Raise(self, node: ast.Raise):
+        for err_node, scope in self._potential_120:
+            self.error(err_node, scope, error_code="ASYNC120")
+        self._potential_120.clear()
 
     def is_safe_aclose_call(self, node: ast.Await) -> bool:
         return (
@@ -120,16 +142,21 @@ def visit_Try(self, node: ast.Try):
         self.visit_nodes(node.finalbody)
 
     def visit_ExceptHandler(self, node: ast.ExceptHandler):
-        if self.cancelled_caught:
-            return
-        res = critical_except(node)
-        if res is None:
+        # 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":
             return
 
-        self.save_state(node, "_critical_scope", "_trio_context_managers")
-        self.cancelled_caught = True
+        self.save_state(
+            node, "_critical_scope", "_trio_context_managers", "_potential_120"
+        )
         self._trio_context_managers = []
-        self._critical_scope = res
+        self._potential_120 = []
+
+        if self.cancelled_caught or (res := critical_except(node)) is None:
+            self._critical_scope = Statement("except", node.lineno, node.col_offset)
+        else:
+            self._critical_scope = res
+            self.cancelled_caught = True
 
     def visit_Assign(self, node: ast.Assign):
         # checks for <scopename>.shield = [True/False]
@@ -145,3 +172,24 @@ def visit_Assign(self, node: ast.Assign):
                 and isinstance(node.value, ast.Constant)
             ):
                 last_scope.shielded = node.value.value
+
+    def visit_FunctionDef(
+        self, node: ast.FunctionDef | ast.AsyncFunctionDef | ast.Lambda
+    ):
+        self.save_state(
+            node,
+            "_critical_scope",
+            "_trio_context_managers",
+            "_potential_120",
+            "cancelled_caught",
+        )
+        self._critical_scope = None
+        self._trio_context_managers = []
+        self.cancelled_caught = False
+
+        self._potential_120 = []
+
+    visit_AsyncFunctionDef = visit_FunctionDef
+    # lambda can't contain await, try, except, raise, with, or assignments.
+    # You also can't do assignment expressions with attributes. So we don't need to
+    # do any special handling for them.

tests/eval_files/async102.py

@@ -1,4 +1,5 @@
 # type: ignore
+# ARG --enable=ASYNC102,ASYNC120
 # NOASYNCIO # TODO: support asyncio shields
 from contextlib import asynccontextmanager
 
@@ -180,7 +181,7 @@ async def foo4():
     try:
         ...
     except ValueError:
-        await foo()  # safe
+        await foo()
     except BaseException:
         await foo()  # error: 8, Statement("BaseException", lineno-1)
     except:
@@ -247,7 +248,7 @@ async def foo_nested_excepts():
         try:
             await foo()  # error: 12, Statement("BaseException", lineno-3)
         except BaseException:
-            await foo()  # error: 12, Statement("BaseException", lineno-1)
+            await foo()  # error: 12, Statement("BaseException", lineno-5)
         except:
             # unsafe, because we're waiting inside the parent baseexception
             await foo()  # error: 12, Statement("BaseException", lineno-8)
@@ -275,3 +276,13 @@ async def foo_nested_async_for():
                 j
             ) in trio.bypasslinters:
                 ...
+
+
+# nested funcdef (previously false alarm)
+async def foo_nested_funcdef():
+    try:
+        ...
+    except:
+
+        async def foobar():
+            await foo()

tests/eval_files/async102_aclose.py

@@ -1,8 +1,12 @@
 # type: ignore
+# ARG --enable=ASYNC102,ASYNC120
 
-# exclude finally: await x.aclose() from async102, with trio/anyio
+# exclude finally: await x.aclose() from async102 and async120, with trio/anyio
+
+# These magical markers are the ones that ensure trio & anyio don't raise errors:
 # ANYIO_NO_ERROR
 # TRIO_NO_ERROR
+
 # See also async102_aclose_args.py - which makes sure trio/anyio raises errors if there
 # are arguments to aclose().
 

tests/eval_files/async120.py

@@ -0,0 +1,123 @@
+# ARG --enable=ASYNC102,ASYNC120
+# NOASYNCIO # TODO: support asyncio shields (?)
+
+import trio
+
+
+def condition() -> bool:
+    return False
+
+
+async def foo():
+    try:
+        ...
+    except ValueError:
+        await foo()  # ASYNC120: 8, Stmt("except", lineno-1)
+        if condition():
+            raise
+        await foo()
+
+    try:
+        ...
+    except ValueError:
+        await foo()  # ASYNC120: 8, Stmt("except", lineno-1)
+        raise
+
+    # don't error if the raise is in a separate excepthandler
+    try:
+        ...
+    except ValueError:
+        await foo()
+    except TypeError:
+        raise
+
+    # does not support conditional branches
+    try:
+        ...
+    except ValueError:
+        if ...:
+            await foo()  # ASYNC120: 12, Stmt("except", lineno-2)
+        else:
+            raise
+
+    # don't trigger on cases of ASYNC102 (?)
+    try:
+        ...
+    except:
+        await foo()  # ASYNC102: 8, Stmt("bare except", lineno-1)
+        raise
+
+    # shielded awaits with timeouts don't trigger 120
+    try:
+        ...
+    except:
+        with trio.fail_after(10) as cs:
+            cs.shield = True
+            await foo()
+        raise
+
+    try:
+        ...
+    except:
+        with trio.fail_after(10) as cs:
+            cs.shield = True
+            await foo()
+            raise
+
+    # ************************
+    # Weird nesting edge cases
+    # ************************
+
+    # nested excepthandlers should not trigger 120 on awaits in
+    # their parent scope
+    try:
+        ...
+    except ValueError:
+        await foo()
+        try:
+            ...
+        except TypeError:
+            raise
+
+    # but the other way around probably should(?)
+    try:
+        ...
+    except ValueError:
+        try:
+            ...
+        except TypeError:
+            await foo()
+        raise
+
+    # but only when they're properly nested, this should not give 120
+    try:
+        ...
+    except TypeError:
+        await foo()
+    if condition():
+        raise
+
+    try:
+        ...
+    except ValueError:
+        await foo()  # ASYNC120: 8, Statement("except", lineno-1)
+        try:
+            await foo()  # ASYNC120: 12, Statement("except", lineno-3)
+        except BaseException:
+            await foo()  # ASYNC102: 12, Statement("BaseException", lineno-1)
+        except:
+            await foo()
+        await foo()  # ASYNC120: 8, Statement("except", lineno-8)
+        raise
+
+
+# nested funcdef
+async def foo_nested_funcdef():
+    try:
+        ...
+    except ValueError:
+
+        async def foobar():
+            await foo()
+
+        raise