fix various issues with ASYNC102 (#289)

- ASYNC102 and ASYNC120 now
  - handles nested cancel scopes
  - detects internal cancel scopes of nurseries as a way to shield&deadline
  - no longer treats trio.open_nursery or anyio.create_task_group as cancellation sources
  - handles the `shield` parameter to trio.fail_after and friends

Author: jakkdl

docs/changelog.rst

@@ -4,6 +4,14 @@ Changelog
 
 `CalVer, YY.month.patch <https://calver.org/>`_
 
+24.9.3
+======
+- :ref:`ASYNC102 <async102>` and :ref:`ASYNC120 <async120>`:
+  - handles nested cancel scopes
+  - detects internal cancel scopes of nurseries as a way to shield&deadline
+  - no longer treats :func:`trio.open_nursery` or :func:`anyio.create_task_group` as cancellation sources
+  - handles the `shield` parameter to :func:`trio.fail_after` and friends (added in trio 0.27)
+
 24.9.2
 ======
 - Fix false alarm in :ref:`ASYNC113 <async113>` and :ref:`ASYNC121 <async121>` with sync functions nested inside an async function.

flake8_async/__init__.py

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

flake8_async/visitors/visitor102.py

@@ -42,13 +42,17 @@ def __init__(self, node: ast.Call, funcname: str, _):
 
             if self.funcname == "CancelScope":
                 self.has_timeout = False
+                for kw in node.keywords:
+                    # note: sets to True even if timeout is explicitly set to inf
+                    if kw.arg == "deadline":
+                        self.has_timeout = True
+
+            # trio 0.27 adds shield parameter to all scope helpers
+            if self.funcname in cancel_scope_names:
                 for kw in node.keywords:
                     # Only accepts constant values
                     if kw.arg == "shield" and isinstance(kw.value, ast.Constant):
                         self.shielded = kw.value.value
-                    # sets to True even if timeout is explicitly set to inf
-                    if kw.arg == "deadline":
-                        self.has_timeout = True
 
     def __init__(self, *args: Any, **kwargs: Any):
         super().__init__(*args, **kwargs)
@@ -109,7 +113,12 @@ def visit_With(self, node: ast.With | ast.AsyncWith):
 
         # Check for a `with trio.<scope_creator>`
         for item in node.items:
-            call = get_matching_call(item.context_expr, *cancel_scope_names)
+            call = get_matching_call(
+                item.context_expr,
+                "open_nursery",
+                "create_task_group",
+                *cancel_scope_names,
+            )
             if call is None:
                 continue
 
@@ -122,7 +131,18 @@ def visit_With(self, node: ast.With | ast.AsyncWith):
             break
 
     def visit_AsyncWith(self, node: ast.AsyncWith):
-        self.async_call_checker(node)
+        # trio.open_nursery and anyio.create_task_group are not cancellation points
+        # so only treat this as an async call if it contains a call that does not match.
+        # asyncio.TaskGroup() appears to be a source of cancellation when exiting.
+        for item in node.items:
+            if not (
+                get_matching_call(item.context_expr, "open_nursery", base="trio")
+                or get_matching_call(
+                    item.context_expr, "create_task_group", base="anyio"
+                )
+            ):
+                self.async_call_checker(node)
+                break
         self.visit_With(node)
 
     def visit_Try(self, node: ast.Try):
@@ -160,18 +180,31 @@ def visit_ExceptHandler(self, node: ast.ExceptHandler):
 
     def visit_Assign(self, node: ast.Assign):
         # checks for <scopename>.shield = [True/False]
+        # and <scopename>.cancel_scope.shield
+        # We don't care to differentiate between them depending on if the scope is
+        # a nursery or not, so e.g. `cs.cancel_scope.shield`/`nursery.shield` will "work"
         if self._trio_context_managers and len(node.targets) == 1:
-            last_scope = self._trio_context_managers[-1]
             target = node.targets[0]
-            if (
-                last_scope.variable_name is not None
-                and isinstance(target, ast.Attribute)
-                and isinstance(target.value, ast.Name)
-                and target.value.id == last_scope.variable_name
-                and target.attr == "shield"
-                and isinstance(node.value, ast.Constant)
-            ):
-                last_scope.shielded = node.value.value
+            for scope in reversed(self._trio_context_managers):
+                if (
+                    scope.variable_name is not None
+                    and isinstance(node.value, ast.Constant)
+                    and isinstance(target, ast.Attribute)
+                    and target.attr == "shield"
+                    and (
+                        (
+                            isinstance(target.value, ast.Name)
+                            and target.value.id == scope.variable_name
+                        )
+                        or (
+                            isinstance(target.value, ast.Attribute)
+                            and target.value.attr == "cancel_scope"
+                            and isinstance(target.value.value, ast.Name)
+                            and target.value.value.id == scope.variable_name
+                        )
+                    )
+                ):
+                    scope.shielded = node.value.value
 
     def visit_FunctionDef(
         self, node: ast.FunctionDef | ast.AsyncFunctionDef | ast.Lambda

tests/eval_files/async102.py

@@ -21,6 +21,12 @@ async def foo():
             s.shield = True
             await foo()
 
+    try:
+        pass
+    finally:
+        with trio.move_on_after(30, shield=True) as s:
+            await foo()
+
     try:
         pass
     finally:
@@ -116,15 +122,6 @@ async def foo():
             await foo()  # error: 12, Statement("try/finally", lineno-7)
     try:
         pass
-    finally:
-        # false alarm, open_nursery does not block/checkpoint on entry.
-        async with trio.open_nursery() as nursery:  # error: 8, Statement("try/finally", lineno-4)
-            nursery.cancel_scope.deadline = trio.current_time() + 10
-            nursery.cancel_scope.shield = True
-            # false alarm, we currently don't handle nursery.cancel_scope.[deadline/shield]
-            await foo()  # error: 12, Statement("try/finally", lineno-8)
-    try:
-        pass
     finally:
         with trio.CancelScope(deadline=30, shield=True):
             with trio.move_on_after(30):
@@ -286,3 +283,24 @@ async def foo_nested_funcdef():
 
         async def foobar():
             await foo()
+
+
+# nested cs
+async def foo_nested_cs():
+    try:
+        ...
+    except:
+        with trio.CancelScope(deadline=10) as cs1:
+            with trio.CancelScope(deadline=10) as cs2:
+                await foo()  # error: 16, Statement("bare except", lineno-3)
+                cs1.shield = True
+                await foo()
+                cs1.shield = False
+                await foo()  # error: 16, Statement("bare except", lineno-7)
+                cs2.shield = True
+                await foo()
+            await foo()  # error: 12, Statement("bare except", lineno-10)
+            cs2.shield = True
+            await foo()  # error: 12, Statement("bare except", lineno-12)
+            cs1.shield = True
+            await foo()

tests/eval_files/async102_anyio.py

@@ -72,3 +72,16 @@ async def foo_anyio_shielded():
             await foo()  # safe
     except BaseException:
         await foo()  # safe
+
+
+# anyio.create_task_group is not a source of cancellations
+async def foo_open_nursery_no_cancel():
+    try:
+        pass
+    finally:
+        # create_task_group does not block/checkpoint on entry, and is not
+        # a cancellation point on exit.
+        async with anyio.create_task_group() as tg:
+            tg.cancel_scope.deadline = anyio.current_time() + 10
+            tg.cancel_scope.shield = True
+            await foo()

tests/eval_files/async102_asyncio.py

@@ -37,3 +37,12 @@ async def foo():
         await asyncio.shield(  # error: 8, Statement("try/finally", lineno-3)
             asyncio.wait_for(foo())
         )
+
+
+# asyncio.TaskGroup *is* a source of cancellations (on exit)
+async def foo_open_nursery_no_cancel():
+    try:
+        pass
+    finally:
+        async with asyncio.TaskGroup() as tg:  # error: 8, Statement("try/finally", lineno-3)
+            ...

tests/eval_files/async102_trio.py

@@ -31,3 +31,16 @@ async def foo5():
             await foo()  # safe
     except BaseException:
         await foo()  # safe, since after trio.Cancelled
+
+
+# trio.open_nursery is not a source of cancellations
+async def foo_open_nursery_no_cancel():
+    try:
+        pass
+    finally:
+        # open_nursery does not block/checkpoint on entry, and is not
+        # a cancellation point on exit.
+        async with trio.open_nursery() as nursery:
+            nursery.cancel_scope.deadline = trio.current_time() + 10
+            nursery.cancel_scope.shield = True
+            await foo()