Trio101: Never yield inside a nursery or cancel scope

README.md

@@ -27,3 +27,4 @@ pip install git+https://github.com/Zac-HD/flake8-trio
 - **TRIO100**: a `with trio.fail_after(...):` or `with trio.move_on_after(...):`
   context does not contain any `await` statements.  This makes it pointless, as
   the timeout can only be triggered by a checkpoint.
+- **TRIO101** `yield` inside a nursery or cancel scope is only safe when implementing a context manager - otherwise, it breaks exception handling.

flake8_trio.py

@@ -11,13 +11,20 @@
 
 import ast
 import tokenize
-from typing import Any, Generator, List, Optional, Tuple, Type, Union
+from typing import Any, Generator, List, Optional, Set, Tuple, Type, Union
 
 # CalVer: YY.month.patch, e.g. first release of July 2022 == "22.7.1"
 __version__ = "22.7.1"
 
 
 Error = Tuple[int, int, str, Type[Any]]
+cancel_scope_names = (
+    "fail_after",
+    "fail_at",
+    "move_on_after",
+    "move_at",
+    "CancelScope",
+)
 
 
 def make_error(error: str, lineno: int, col: int, *args: Any, **kwargs: Any) -> Error:
@@ -40,19 +47,58 @@ class Visitor(ast.NodeVisitor):
     def __init__(self) -> None:
         super().__init__()
         self.problems: List[Error] = []
+        self.safe_yields: Set[ast.Yield] = set()
+        self._yield_is_error = False
+        self._context_manager = False
 
-    def visit_With(self, node: ast.With) -> None:
+    def visit_generic_with(self, node: Union[ast.With, ast.AsyncWith]):
         self.check_for_trio100(node)
+
+        outer = self._yield_is_error
+        if not self._context_manager and any(
+            is_trio_call(item, "open_nursery", *cancel_scope_names)
+            for item in (i.context_expr for i in node.items)
+        ):
+            self._yield_is_error = True
+
         self.generic_visit(node)
+        self._yield_is_error = outer
+
+    def visit_With(self, node: ast.With) -> None:
+        self.visit_generic_with(node)
 
     def visit_AsyncWith(self, node: ast.AsyncWith) -> None:
-        self.check_for_trio100(node)
+        self.visit_generic_with(node)
+
+    def visit_generic_FunctionDef(
+        self, node: Union[ast.FunctionDef, ast.AsyncFunctionDef]
+    ):
+        outer = self._context_manager
+        if any(
+            isinstance(d, ast.Name)
+            and d.id in ("contextmanager", "asynccontextmanager")
+            for d in node.decorator_list
+        ):
+            self._context_manager = True
+        self.generic_visit(node)
+        self._context_manager = outer
+
+    def visit_FunctionDef(self, node: ast.FunctionDef) -> None:
+        self.visit_generic_FunctionDef(node)
+
+    def visit_AsyncFunctionDef(self, node: ast.AsyncFunctionDef) -> None:
+        self.visit_generic_FunctionDef(node)
+
+    def visit_Yield(self, node: ast.Yield) -> None:
+        if self._yield_is_error:
+            self.problems.append(make_error(TRIO101, node.lineno, node.col_offset))
+
         self.generic_visit(node)
 
     def check_for_trio100(self, node: Union[ast.With, ast.AsyncWith]) -> None:
         # Context manager with no `await` call within
         for item in (i.context_expr for i in node.items):
-            call = is_trio_call(item, "fail_after", "move_on_after")
+            call = is_trio_call(item, *cancel_scope_names)
             if call and not any(isinstance(x, ast.Await) for x in ast.walk(node)):
                 self.problems.append(
                     make_error(TRIO100, item.lineno, item.col_offset, call)
@@ -79,3 +125,4 @@ def run(self) -> Generator[Tuple[int, int, str, Type[Any]], None, None]:
 
 
 TRIO100 = "TRIO100: {} context contains no checkpoints, add `await trio.sleep(0)`"
+TRIO101 = "TRIO101: yield inside a nursery or cancel scope is only safe when implementing a context manager - otherwise, it breaks exception handling"

tests/test_flake8_trio.py

@@ -9,7 +9,7 @@
 from hypothesis import HealthCheck, given, settings
 from hypothesmith import from_grammar, from_node
 
-from flake8_trio import TRIO100, Error, Plugin, Visitor, make_error
+from flake8_trio import TRIO100, TRIO101, Error, Plugin, Visitor, make_error
 
 
 class Flake8TrioTestCase(unittest.TestCase):
@@ -40,6 +40,15 @@ def test_trio100_py39(self):
             make_error(TRIO100, 14, 8, "trio.move_on_after"),
         )
 
+    def test_trio101(self):
+        self.maxDiff = None
+        self.assert_expected_errors(
+            "trio101.py",
+            make_error(TRIO101, 8, 8),
+            make_error(TRIO101, 13, 8),
+            make_error(TRIO101, 25, 8),
+        )
+
 
 @pytest.mark.fuzz
 class TestFuzz(unittest.TestCase):

tests/trio101.py

@@ -0,0 +1,31 @@
+from contextlib import asynccontextmanager, contextmanager
+
+import trio
+
+
+def foo0():
+    with trio.open_nursery() as _:
+        yield 1  # error
+
+
+async def foo1():
+    async with trio.open_nursery() as _:
+        yield 1  # error
+
+
+@contextmanager
+def foo2():
+    with trio.open_nursery() as _:
+        yield 1  # safe
+
+
+async def foo3():
+    async with trio.CancelScope() as _:
+        await trio.sleep(1)  # so trio100 doesn't complain
+        yield 1  # error
+
+
+@asynccontextmanager
+async def foo4():
+    async with trio.open_nursery() as _:
+        yield 1  # safe