ASYNC100 now treats start_soon() as a cancel point (#327)

* ASYNC100 now treats start_soon() as a cancel point

* oh derp, visitor91x is a CST visitor

* start_soon now makes open_nursery a cancel point on exit

docs/changelog.rst

@@ -4,9 +4,13 @@ Changelog
 
 `CalVer, YY.month.patch <https://calver.org/>`_
 
+24.11.4
+=======
+- :ref:`ASYNC100 <async100>` once again ignores :func:`trio.open_nursery` and :func:`anyio.create_task_group`, unless we find a call to ``.start_soon()``.
+
 24.11.3
 =======
-- Revert :ref:`ASYNC100 <async100>` ignoring :func:`trio.open_nursery` and :func:`anyio.create_task_group` due to it not viewing `start_soon()` as introducing a :ref:`cancel point <cancel_point>`.
+- Revert :ref:`ASYNC100 <async100>` ignoring :func:`trio.open_nursery` and :func:`anyio.create_task_group` due to it not viewing ``.start_soon()`` as introducing a :ref:`cancel point <cancel_point>`.
 
 24.11.2
 =======

docs/glossary.rst

@@ -99,7 +99,7 @@ functions defined by Trio will either checkpoint or raise an exception when
 iteration, and when exhausting the iterator, and ``async with`` will checkpoint
 on at least one of enter/exit.
 
-The one exception is :func:`trio.open_nursery` and :func:`anyio.create_task_group` which are :ref:`schedule points <schedule_point>` but not :ref:`cancel points <cancel_point>`.
+The one exception is :func:`trio.open_nursery` and :func:`anyio.create_task_group`. They do not checkpoint on entry, and on exit they insert a :ref:`schedule point <schedule_point>`. However, if sub-tasks are cancelled they will be propagated on exit, so if you're starting tasks you can usually treat the exit as a :ref:`cancel point <cancel_point>`.
 
 asyncio does not place any guarantees on if or when asyncio functions will
 checkpoint. This means that enabling and adhering to :ref:`ASYNC91x <ASYNC910>`

docs/rules.rst

@@ -13,6 +13,7 @@ _`ASYNC100` : cancel-scope-no-checkpoint
     A :ref:`timeout_context` does not contain any :ref:`checkpoints <checkpoint>`.
     This makes it pointless, as the timeout can only be triggered by a checkpoint.
     This check also treats ``yield`` as a checkpoint, since checkpoints can happen in the caller we yield to.
+    :func:`trio.open_nursery` and :func:`anyio.create_task_group` are excluded, as they are :ref:`schedule points <schedule_point>` but not :ref:`cancel points <cancel_point>` (unless they have tasks started in them).
     See :ref:`ASYNC912 <async912>` which will in addition guarantee checkpoints on every code path.
 
 _`ASYNC101` : yield-in-cancel-scope

flake8_async/__init__.py

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

flake8_async/visitors/visitor91x.py

@@ -25,6 +25,7 @@
     flatten_preserving_comments,
     fnmatch_qualified_name_cst,
     func_has_decorator,
+    identifier_to_string,
     iter_guaranteed_once_cst,
     with_has_call,
 )
@@ -285,6 +286,7 @@ def __init__(self, *args: Any, **kwargs: Any):
         # ASYNC100
         self.has_checkpoint_stack: list[bool] = []
         self.node_dict: dict[cst.With, list[AttributeCall]] = {}
+        self.taskgroup_has_start_soon: dict[str, bool] = {}
 
         # --exception-suppress-context-manager
         self.suppress_imported_as: list[str] = []
@@ -299,13 +301,31 @@ def should_autofix(self, node: cst.CSTNode, code: str | None = None) -> bool:
             and self.library != ("asyncio",)
         )
 
-    def checkpoint(self) -> None:
-        self.uncheckpointed_statements = set()
+    def checkpoint_cancel_point(self) -> None:
         self.has_checkpoint_stack = [True] * len(self.has_checkpoint_stack)
+        # don't need to look for any .start_soon() calls
+        self.taskgroup_has_start_soon.clear()
+
+    def checkpoint_schedule_point(self) -> None:
+        self.uncheckpointed_statements = set()
+
+    def checkpoint(self) -> None:
+        self.checkpoint_schedule_point()
+        self.checkpoint_cancel_point()
 
     def checkpoint_statement(self) -> cst.SimpleStatementLine:
         return checkpoint_statement(self.library[0])
 
+    def visit_Call(self, node: cst.Call) -> None:
+        # [Nursery/TaskGroup].start_soon introduces a cancel point
+        if (
+            isinstance(node.func, cst.Attribute)
+            and isinstance(node.func.value, cst.Name)
+            and node.func.attr.value == "start_soon"
+            and node.func.value.value in self.taskgroup_has_start_soon
+        ):
+            self.taskgroup_has_start_soon[node.func.value.value] = True
+
     def visit_ImportFrom(self, node: cst.ImportFrom) -> None:
         # Semi-crude approach to handle `from contextlib import suppress`.
         # It does not handle the identifier being overridden, or assigned
@@ -341,16 +361,21 @@ def visit_FunctionDef(self, node: cst.FunctionDef) -> bool:
             "safe_decorator",
             "async_function",
             "uncheckpointed_statements",
+            # comp_unknown does not need to be saved
             "loop_state",
             "try_state",
             "has_checkpoint_stack",
-            "suppress_imported_as",
+            # node_dict is cleaned up and don't need to be saved
+            "taskgroup_has_start_soon",
+            "suppress_imported_as",  # a copy is saved, but state is not reset
             copy=True,
         )
-        self.uncheckpointed_statements = set()
-        self.has_checkpoint_stack = []
         self.has_yield = self.safe_decorator = False
+        self.uncheckpointed_statements = set()
         self.loop_state = LoopState()
+        # try_state is reset upon entering try
+        self.has_checkpoint_stack = []
+        self.taskgroup_has_start_soon = {}
 
         self.async_function = (
             node.asynchronous is not None
@@ -440,8 +465,8 @@ def leave_Return(
         ):
             self.add_statement = self.checkpoint_statement()
         # avoid duplicate error messages
-        self.uncheckpointed_statements = set()
-        # we don't treat it as a checkpoint for ASYNC100
+        # but don't see it as a cancel point for ASYNC100
+        self.checkpoint_schedule_point()
 
         # return original node to avoid problems with identity equality
         assert original_node.deep_equals(updated_node)
@@ -491,12 +516,47 @@ def _is_exception_suppressing_context_manager(self, node: cst.With) -> bool:
             is not None
         )
 
+    def _checkpoint_with(self, node: cst.With, entry: bool):
+        """Conditionally checkpoints entry/exit of With.
+
+        If the `with` only contains calls to open_nursery/create_task_group, it's a
+        schedule point but not a cancellation point, so we treat it as a checkpoint
+        for async91x but not for async100.
+
+        Saves the name of the taskgroup/nursery if entry is set
+        """
+        if not getattr(node, "asynchronous", None):
+            return
+
+        for item in node.items:
+            if isinstance(item.item, cst.Call) and identifier_to_string(
+                item.item.func
+            ) in (
+                "trio.open_nursery",
+                "anyio.create_task_group",
+            ):
+                if item.asname is not None and isinstance(item.asname.name, cst.Name):
+                    # save the nursery/taskgroup to see if it has a `.start_soon`
+                    if entry:
+                        self.taskgroup_has_start_soon[item.asname.name.value] = False
+                    elif self.taskgroup_has_start_soon.pop(
+                        item.asname.name.value, False
+                    ):
+                        self.checkpoint()
+                        return
+            else:
+                self.checkpoint()
+                break
+        else:
+            # only taskgroup/nursery calls
+            self.checkpoint_schedule_point()
+
     # Async context managers can reasonably checkpoint on either or both of entry and
     # exit.  Given that we can't tell which, we assume "both" to avoid raising a
     # missing-checkpoint warning when there might in fact be one (i.e. a false alarm).
     def visit_With_body(self, node: cst.With):
-        if getattr(node, "asynchronous", None):
-            self.checkpoint()
+        self.save_state(node, "taskgroup_has_start_soon", copy=True)
+        self._checkpoint_with(node, entry=True)
 
         # if this might suppress exceptions, we cannot treat anything inside it as
         # checkpointing.
@@ -548,15 +608,16 @@ def leave_With(self, original_node: cst.With, updated_node: cst.With):
                 for res in self.node_dict[original_node]:
                     self.error(res.node, error_code="ASYNC912")
 
+        self.node_dict.pop(original_node, None)
+
         # if exception-suppressing, restore all uncheckpointed statements from
         # before the `with`.
         if self._is_exception_suppressing_context_manager(original_node):
             prev_checkpoints = self.uncheckpointed_statements
             self.restore_state(original_node)
             self.uncheckpointed_statements.update(prev_checkpoints)
 
-        if getattr(original_node, "asynchronous", None):
-            self.checkpoint()
+        self._checkpoint_with(original_node, entry=False)
         return updated_node
 
     # error if no checkpoint since earlier yield or function entry
@@ -569,7 +630,7 @@ def leave_Yield(
 
         # Treat as a checkpoint for ASYNC100, since the context we yield to
         # may checkpoint.
-        self.has_checkpoint_stack = [True] * len(self.has_checkpoint_stack)
+        self.checkpoint_cancel_point()
 
         if self.check_function_exit(original_node) and self.should_autofix(
             original_node

tests/autofix_files/async100.py

@@ -132,12 +132,6 @@ async def fn(timeout):
             await trio.sleep(1)
 
 
-async def nursery_no_cancel_point():
-    with trio.CancelScope():  # should error, but reverted PR
-        async with anyio.create_task_group():
-            ...
-
-
 async def dont_crash_on_non_name_or_attr_call():
     async with contextlib.asynccontextmanager(agen_fn)():
         ...

tests/autofix_files/async100_anyio.py

@@ -0,0 +1,26 @@
+# AUTOFIX
+# BASE_LIBRARY anyio
+# NOTRIO  # trio.create_task_group doesn't exist
+# ASYNCIO_NO_ERROR
+import anyio
+
+
+async def bar() -> None: ...
+
+
+async def anyio_cancelscope():
+    # error: 9, "anyio", "CancelScope"
+    ...
+
+
+# see async100_trio for more comprehensive tests
+async def nursery_no_cancel_point():
+    # error: 9, "anyio", "CancelScope"
+    async with anyio.create_task_group():
+        ...
+
+
+async def nursery_with_start_soon():
+    with anyio.CancelScope():
+        async with anyio.create_task_group() as tg:
+            tg.start_soon(bar)

tests/autofix_files/async100_anyio.py.diff

@@ -0,0 +1,23 @@
+---
++++
+@@ x,15 x,15 @@
+
+
+ async def anyio_cancelscope():
+-    with anyio.CancelScope():  # error: 9, "anyio", "CancelScope"
+-        ...
++    # error: 9, "anyio", "CancelScope"
++    ...
+
+
+ # see async100_trio for more comprehensive tests
+ async def nursery_no_cancel_point():
+-    with anyio.CancelScope():  # error: 9, "anyio", "CancelScope"
+-        async with anyio.create_task_group():
+-            ...
++    # error: 9, "anyio", "CancelScope"
++    async with anyio.create_task_group():
++        ...
+
+
+ async def nursery_with_start_soon():

tests/autofix_files/async100_trio.py

@@ -1,10 +1,75 @@
 # AUTOFIX
 # ASYNCIO_NO_ERROR # asyncio.open_nursery doesn't exist
-# ANYIO_NO_ERROR # anyio.open_nursery doesn't exist
+# NOANYIO # anyio.open_nursery doesn't exist
 import trio
 
 
 async def nursery_no_cancel_point():
-    with trio.CancelScope():  # should error, but reverted PR
-        async with trio.open_nursery():
-            ...
+    # error: 9, "trio", "CancelScope"
+    async with trio.open_nursery():
+        ...
+
+
+# but it is a cancel point if the nursery contains a call to start_soon()
+
+
+async def nursery_start_soon():
+    with trio.CancelScope():
+        async with trio.open_nursery() as n:
+            n.start_soon(trio.sleep, 0)
+
+
+async def nursery_start_soon_misnested():
+    async with trio.open_nursery() as n:
+        # error: 13, "trio", "CancelScope"
+        n.start_soon(trio.sleep, 0)
+
+
+async def nested_scope():
+    with trio.CancelScope():
+        with trio.CancelScope():
+            async with trio.open_nursery() as n:
+                n.start_soon(trio.sleep, 0)
+
+
+async def nested_nursery():
+    with trio.CancelScope():
+        async with trio.open_nursery() as n:
+            async with trio.open_nursery() as n2:
+                n2.start_soon(trio.sleep, 0)
+
+
+async def nested_function_call():
+
+    # error: 9, "trio", "CancelScope"
+    async with trio.open_nursery() as n:
+
+        def foo():
+            n.start_soon(trio.sleep, 0)
+
+        # a false alarm in case we call foo()... but we can't check if they do
+        foo()
+
+
+# insert cancel point on nursery exit, not at the start_soon call
+async def cancel_point_on_nursery_exit():
+    with trio.CancelScope():
+        async with trio.open_nursery() as n:
+            # error: 17, "trio", "CancelScope"
+            n.start_soon(trio.sleep, 0)
+
+
+# async100 does not consider *redundant* cancel scopes
+async def redundant_cancel_scope():
+    with trio.CancelScope():
+        with trio.CancelScope():
+            await trio.lowlevel.checkpoint()
+
+
+# but if it did then none of these scopes should be marked redundant
+# The inner checks task startup, the outer checks task exit
+async def nursery_exit_blocks_with_start():
+    with trio.CancelScope():
+        async with trio.open_nursery() as n:
+            with trio.CancelScope():
+                await n.start(trio.sleep, 0)