Merge pull request #60 from jakkdl/prefer_start_in_cm

Fixes #57

Author: Zac-HD

CHANGELOG.md

@@ -1,6 +1,9 @@
 # Changelog
 *[CalVer, YY.month.patch](https://calver.org/)*
 
+## 22.11.3
+- Add TRIO113, prefer `await nursery.start(...)` to `nursery.start_soon()` for compatible functions when opening a context manager
+
 ## 22.11.2
 - TRIO105 now also checks that you `await`ed `nursery.start()`.
 

README.md

@@ -35,15 +35,34 @@ pip install flake8-trio
 - **TRIO110**: `while <condition>: await trio.sleep()` should be replaced by a `trio.Event`.
 - **TRIO111**: Variable, from context manager opened inside nursery, passed to `start[_soon]` might be invalidly accesed while in use, due to context manager closing before the nursery. This is usually a bug, and nurseries should generally be the inner-most context manager.
 - **TRIO112**: nursery body with only a call to `nursery.start[_soon]` and not passing itself as a parameter can be replaced with a regular function call.
+- **TRIO113**: using `nursery.start_soon` in `__aenter__` doesn't wait for the task to begin. Consider replacing with `nursery.start`.
 
 
 ## Configuration
-`no-checkpoint-warning-decorators`: Specify a list of decorators to disable checkpointing checks for, turning off TRIO107 and TRIO108 warnings for functions decorated with any decorator matching any in the list. Matching is done with [fnmatch](https://docs.python.org/3/library/fnmatch.html). Defaults to disabling for `asynccontextmanager`.
+### `no-checkpoint-warning-decorators`
+Specify a list of decorators to disable checkpointing checks for, turning off TRIO107 and TRIO108 warnings for functions decorated with any decorator matching any in the list. Matching is done with [fnmatch](https://docs.python.org/3/library/fnmatch.html). Defaults to disabling for `asynccontextmanager`.
 
 Decorators-to-match must be identifiers or dotted names only (not PEP-614 expressions), and will match against the name only - e.g. `foo.bar` matches `foo.bar`, `foo.bar()`, and `foo.bar(args, here)`, etc.
 
-For example:
+[You can configure `flake8` with command-line options](https://flake8.pycqa.org/en/latest/user/configuration.html),
+but we prefer using a config file. For example:
 ```
 [flake8]
-no-checkpoint-warning-decorators = mydecorator, mydecoratorpackage.checkpointing_decorators.*, ign*, *.ignore
+no-checkpoint-warning-decorators =
+  mydecorator,
+  mydecoratorpackage.checkpointing_decorators.*,
+  ign*,
+  *.ignore,
+```
+
+
+### `startable-in-context-manager`
+Comma-separated list of methods which should be used with `.start()` when opening a context manager,
+in addition to the default `trio.run_process`, `trio.serve_tcp`, `trio.serve_ssl_over_tcp`, and
+`trio.serve_listeners`.  Uses fnmatch, like `no-checkpoint-warning-decorators`. For example:
+```
+[flake8]
+startable-methods-in-context-manager =
+  mylib.myfun,
+  myfun2,mypackage.myfunlib.*,
 ```

flake8_trio.py

@@ -29,7 +29,7 @@
 from flake8.options.manager import OptionManager
 
 # CalVer: YY.month.patch, e.g. first release of July 2022 == "22.7.1"
-__version__ = "22.11.2"
+__version__ = "22.11.3"
 
 
 Error_codes = {
@@ -65,6 +65,7 @@
         "Nurseries should generally be the inner-most context manager."
     ),
     "TRIO112": "Redundant nursery {}, consider replacing with directly awaiting the function call",
+    "TRIO113": "Dangerous `.start_soon()`, process might not run before `__aenter__` exits. Consider replacing with `.start()`.",
 }
 
 
@@ -126,6 +127,10 @@ def __lt__(self, other: Any) -> bool:
     def __eq__(self, other: Any) -> bool:
         return isinstance(other, Error) and self.cmp() == other.cmp()
 
+    def __repr__(self) -> str:
+        trailer = "".join(f", {x!r}" for x in self.args)
+        return f"<{self.code} error at {self.line}:{self.col}{trailer}>"
+
 
 checkpoint_node_types = (ast.Await, ast.AsyncFor, ast.AsyncWith)
 cancel_scope_names = (
@@ -202,8 +207,9 @@ def has_decorator(decorator_list: List[ast.expr], *names: str):
     return False
 
 
-# matches the full decorator name against fnmatch pattern
-def fnmatch_decorator(decorator_list: List[ast.expr], *patterns: str):
+# matches the fully qualified name against fnmatch pattern
+# used to match decorators and methods to user-supplied patterns
+def fnmatch_qualified_name(name_list: List[ast.expr], *patterns: str):
     def construct_name(expr: ast.expr) -> Optional[str]:
         if isinstance(expr, ast.Call):
             expr = expr.func
@@ -216,12 +222,13 @@ def construct_name(expr: ast.expr) -> Optional[str]:
         # See https://peps.python.org/pep-0614/ - we don't handle everything
         return None  # pragma: no cover  # impossible on Python 3.8
 
-    for decorator in decorator_list:
-        qualified_decorator_name = construct_name(decorator)
-        if qualified_decorator_name is None:
+    for name in name_list:
+        qualified_name = construct_name(name)
+        if qualified_name is None:
             continue  # pragma: no cover  # impossible on Python 3.8
         for pattern in patterns:
-            if fnmatch(qualified_decorator_name, pattern.lstrip("@")):
+            # strip leading "@"s for when we're working with decorators
+            if fnmatch(qualified_name, pattern.lstrip("@")):
                 return True
     return False
 
@@ -803,6 +810,16 @@ def iter_guaranteed_once(iterable: ast.expr) -> bool:
 )
 
 
+def is_nursery_call(node: ast.AST, name: str) -> bool:  # pragma: no cover
+    assert name in ("start", "start_soon")
+    if isinstance(node, ast.Attribute):
+        if isinstance(node.value, ast.Name):
+            return node.attr == name and node.value.id.endswith("nursery")
+        if isinstance(node.value, ast.Attribute):
+            return node.attr == name and node.value.attr.endswith("nursery")
+    return False
+
+
 class Visitor105(Flake8TrioVisitor):
     def __init__(self, *args: Any, **kwargs: Any):
         super().__init__(*args, **kwargs)
@@ -820,13 +837,14 @@ def visit_Call(self, node: ast.Call):
             and isinstance(node.func.value, ast.Name)
             and (
                 (node.func.value.id == "trio" and node.func.attr in trio_async_funcs)
-                or (node.func.value.id == "nursery" and node.func.attr == "start")
+                or is_nursery_call(node.func, "start")
             )
             and (
                 len(self.node_stack) < 2
                 or not isinstance(self.node_stack[-2], ast.Await)
             )
         ):
+            assert isinstance(node.func, ast.Attribute)
             self.error("TRIO105", node, node.func.attr)
         self.generic_visit(node)
 
@@ -873,7 +891,7 @@ def visit_AsyncFunctionDef(self, node: ast.AsyncFunctionDef):
         self.set_state(self.default, copy=True)
 
         # disable checks in asynccontextmanagers by saying the function isn't async
-        self.async_function = not fnmatch_decorator(
+        self.async_function = not fnmatch_qualified_name(
             node.decorator_list, *self.options.no_checkpoint_warning_decorators
         )
 
@@ -1149,10 +1167,60 @@ def visit_BoolOp(self, node: ast.BoolOp):
         self.uncheckpointed_statements = worst_case_shortcut
 
 
+def _get_identifier(node: ast.expr) -> str:
+    if isinstance(node, ast.Name):
+        return node.id
+    if isinstance(node, ast.Attribute):
+        return node.attr
+    return ""  # pragma: no cover
+
+
+STARTABLE_CALLS = (
+    "trio.run_process",
+    "trio.serve_ssl_over_tcp",
+    "trio.serve_tcp",
+    "trio.serve_listeners",
+)
+
+
+class Visitor113(Flake8TrioVisitor):
+    def __init__(self, *args: Any, **kwargs: Any):
+        super().__init__(*args, **kwargs)
+        self.async_function = False
+        self.asynccontextmanager = False
+        self.aenter = False
+
+    def visit_AsyncFunctionDef(self, node: ast.AsyncFunctionDef):
+        outer = self.aenter
+
+        self.aenter = node.name == "__aenter__" or any(
+            _get_identifier(d) == "asynccontextmanager" for d in node.decorator_list
+        )
+
+        self.generic_visit(node)
+        self.aenter = outer
+
+    def visit_Yield(self, node: ast.Yield):
+        self.aenter = False
+
+    def visit_Call(self, node: ast.Call) -> None:
+        if (
+            self.aenter
+            and is_nursery_call(node.func, "start_soon")
+            and len(node.args) > 0
+            and fnmatch_qualified_name(
+                node.args[:1],
+                *STARTABLE_CALLS,
+                *self.options.startable_in_context_manager,
+            )
+        ):
+            self.error("TRIO113", node)
+
+
 class Plugin:
     name = __name__
     version = __version__
-    options: Namespace = Namespace(no_checkpoint_warning_decorators=[])
+    options: Namespace = Namespace()
 
     def __init__(self, tree: ast.AST):
         self._tree = tree
@@ -1183,6 +1251,21 @@ def add_options(option_manager: OptionManager):
                 "mydecorator,mypackage.mydecorators.*``"
             ),
         )
+        option_manager.add_option(
+            "--startable-in-context-manager",
+            default="",
+            parse_from_config=True,
+            required=False,
+            comma_separated_list=True,
+            help=(
+                "Comma-separated list of method calls to enable TRIO113"
+                " warnings for."
+                " Use if you want to override the default methods the check is enabled for."
+                " Methods can be dotted or not, as well as support * as a wildcard."
+                " For example, ``--startable-in-context-manager=worker_serve,"
+                "mypackage.sub.*``"
+            ),
+        )
 
     @staticmethod
     def parse_options(options: Namespace):

tests/test_decorator.py

@@ -6,7 +6,7 @@
 import pytest
 from flake8.main.application import Application
 
-from flake8_trio import Error_codes, Plugin, Statement, fnmatch_decorator
+from flake8_trio import Error_codes, Plugin, Statement, fnmatch_qualified_name
 
 
 def dec_list(*decorators: str) -> ast.Module:
@@ -21,7 +21,7 @@ def dec_list(*decorators: str) -> ast.Module:
 def wrap(decorators: Tuple[str, ...], decs2: str) -> bool:
     tree = dec_list(*decorators)
     assert isinstance(tree.body[0], ast.AsyncFunctionDef)
-    return fnmatch_decorator(tree.body[0].decorator_list, decs2)
+    return fnmatch_qualified_name(tree.body[0].decorator_list, decs2)
 
 
 def test_basic():
@@ -87,6 +87,7 @@ def test_pep614():
 def test_plugin():
     tree = dec_list("app.route")
     plugin = Plugin(tree)
+    plugin.options = Namespace(no_checkpoint_warning_decorators=[])
     assert tuple(plugin.run())
 
     plugin.options = Namespace(no_checkpoint_warning_decorators=["app.route"])

tests/test_flake8_trio.py

@@ -11,6 +11,7 @@
 from typing import DefaultDict, Iterable, List, Sequence, Tuple, Type
 
 import pytest
+from flake8.options.manager import OptionManager
 
 # import trio  # type: ignore
 from hypothesis import HealthCheck, given, settings
@@ -157,6 +158,15 @@ def read_file(test_file: str):
 
 
 def assert_expected_errors(plugin: Plugin, include: Iterable[str], *expected: Error):
+    # initialize default option values
+    om = OptionManager(
+        version="",
+        plugin_versions="",
+        parents=[],
+        formatter_names=["default"],  # type: ignore
+    )
+    plugin.add_options(om)
+    plugin.parse_options(om.parse_args(args=[""]))
 
     errors = sorted(e for e in plugin.run() if e.code in include)
     expected_ = sorted(expected)
@@ -353,6 +363,26 @@ def test_107_permutations():
             assert not errors, "# false alarm:\n" + function_str
 
 
+def test_113_options():
+    # check that no errors are given by default
+    plugin = read_file("trio113.py")
+    om = OptionManager(
+        version="",
+        plugin_versions="",
+        parents=[],
+        formatter_names=["default"],  # type: ignore
+    )
+    plugin.add_options(om)
+    plugin.parse_options(om.parse_args(args=["--startable-in-context-manager=''"]))
+    default = {repr(e) for e in plugin.run() if e.code == "TRIO113"}
+
+    # and that the expected errors are given if we empty it and then extend it
+    arg = "--startable-in-context-manager=custom_startable_function"
+    plugin.parse_options(om.parse_args(args=[arg]))
+    errors = {repr(e) for e in plugin.run() if e.code == "TRIO113"} - default
+    assert errors == {repr(Error("TRIO113", 58, 8))}
+
+
 @pytest.mark.fuzz
 class TestFuzz(unittest.TestCase):
     @settings(max_examples=1_000, suppress_health_check=[HealthCheck.too_slow])

tests/trio113.py

@@ -0,0 +1,88 @@
+import contextlib
+import contextlib as arbitrary_import_alias_for_contextlib
+import os
+from contextlib import asynccontextmanager
+
+import trio
+
+
+@asynccontextmanager
+async def run_sampling_profiler():
+    async with trio.open_nursery() as nursery:
+        nursery.start_soon(  # error: 8
+            trio.run_process, ["start-sampling-profiler", str(os.getpid())]
+        )
+        yield
+
+        nursery.start_soon(  # safe, it's after the first yield
+            trio.run_process, ["start-sampling-profiler", str(os.getpid())]
+        )
+
+
+async def not_cm():
+    async with trio.open_nursery() as nursery:
+        nursery.start_soon(
+            trio.run_process, ["start-sampling-profiler", str(os.getpid())]
+        )
+        yield
+
+
+class foo:
+    async def __aenter__(self):
+        async with trio.open_nursery() as nursery:
+            nursery.start_soon(  # error: 12
+                trio.run_process, ["start-sampling-profiler", str(os.getpid())]
+            )
+            yield
+
+
+nursery = anything = custom_startable_function = trio  # type: ignore
+
+
+class foo2:
+    _nursery = nursery
+
+    async def __aenter__(self):
+        nursery.start_soon(trio.run_process)  # error: 8
+        nursery.start_soon(trio.serve_ssl_over_tcp)  # error: 8
+        nursery.start_soon(trio.serve_tcp)  # error: 8
+        nursery.start_soon(trio.serve_listeners)  # error: 8
+
+        # Where does `nursery` come from in __aenter__?  Probably an attribute:
+        self._nursery.start_soon(trio.run_process)  # error: 8
+        self._nursery.start_soon(trio.serve_ssl_over_tcp)  # error: 8
+        self._nursery.start_soon(trio.serve_tcp)  # error: 8
+        self._nursery.start_soon(trio.serve_listeners)  # error: 8
+
+        # We have a test that sets `startable-in-context-manager` to error here
+        nursery.start_soon(custom_startable_function)
+
+
+class foo3:
+    async def __aenter__(_a_parameter_not_named_self):
+        nursery.start_soon(trio.run_process)  # error: 8
+
+
+# might be monkeypatched onto an instance, count this as an error too
+async def __aenter__():
+    nursery.start_soon(trio.run_process)  # error: 4
+    nursery.start_soon()  # broken code, but our analysis shouldn't crash
+    nursery.cancel_scope.cancel()
+
+
+# this only takes a single parameter ... right? :P
+class foo4:
+    async def __aenter__(self, *, args="foo"):
+        nursery.start_soon(trio.run_process)  # error: 8
+
+
+@contextlib.asynccontextmanager
+async def contextlib_acm():
+    nursery.start_soon(trio.run_process)  # error: 4
+    yield
+
+
+@arbitrary_import_alias_for_contextlib.asynccontextmanager
+async def contextlib_import_alias_acm():
+    nursery.start_soon(trio.run_process)  # error: 4
+    yield