Add config option for async900 decorators

Author: Zac-HD

docs/changelog.rst

@@ -4,6 +4,11 @@ Changelog
 
 *[CalVer, YY.month.patch](https://calver.org/)*
 
+24.8.1
+======
+- Add config option ``transform-async-generator-decorators``, to list decorators which
+  suppress :ref:`ASYNC900 <async900>`.
+
 24.6.1
 ======
 - Add :ref:`ASYNC120 <async120>` await-in-except.

docs/conf.py

@@ -34,6 +34,7 @@
     "python": ("https://docs.python.org/3", None),
     "anyio": ("https://anyio.readthedocs.io/en/latest/", None),
     "trio": ("https://trio.readthedocs.io/en/latest/", None),
+    "trio-util": ("https://trio-util.readthedocs.io/en/latest/", None),
     # httpx? Did not seem to work on first try - I think they might not be using
     # sphinx at all, so probably can't link through intersphinx?
     # see https://github.com/encode/httpx/discussions/1220

docs/rules.rst

@@ -161,12 +161,14 @@ Optional rules disabled by default
 Our 9xx rules check for semantics issues, like 1xx rules, but are disabled by default due
 to the higher volume of warnings.  We encourage you to enable them - without guaranteed
 :ref:`checkpoint`\ s timeouts and cancellation can be arbitrarily delayed, and async
-generators are prone to the problems described in :pep:`533`.
+generators are prone to the problems described in :pep:`789` and :pep:`533`.
 
 _`ASYNC900` : unsafe-async-generator
        Async generator without :func:`@asynccontextmanager <contextlib.asynccontextmanager>` not allowed.
        You might want to enable this on a codebase since async generators are inherently unsafe and cleanup logic might not be performed.
-       See `#211 <https://github.com/python-trio/flake8-async/issues/211>`__ and https://discuss.python.org/t/using-exceptiongroup-at-anthropic-experience-report/20888/6 for discussion.
+       See :pep:`789` for control-flow problems, :pep:`533` for delayed cleanup problems.
+       Further decorators can be registered with the ``--transform-async-generator-decorators``
+       config option, e.g. :obj:`trio_util.trio_async_generator`.
 
 _`ASYNC910` : async-function-no-checkpoint
     Exit or ``return`` from async function with no guaranteed :ref:`checkpoint` or exception since function definition.

flake8_async/__init__.py

@@ -38,7 +38,7 @@
 
 
 # CalVer: YY.month.patch, e.g. first release of July 2022 == "22.7.1"
-__version__ = "24.6.1"
+__version__ = "24.8.1"
 
 
 # taken from https://github.com/Zac-HD/shed
@@ -261,6 +261,18 @@ def add_options(option_manager: OptionManager | ArgumentParser):
                 "mydecorator,mypackage.mydecorators.*``"
             ),
         )
+        add_argument(
+            "--transform-async-generator-decorators",
+            default="",
+            required=False,
+            type=comma_separated_list,
+            help=(
+                "Comma-separated list of decorators to disable ASYNC900 warnings for. "
+                "Decorators can be dotted or not, as well as support * as a wildcard. "
+                "For example, ``--transform-async-generator-decorators=fastapi.Depends,"
+                "trio_util.trio_async_generator``"
+            ),
+        )
         add_argument(
             "--exception-suppress-context-managers",
             default="",
@@ -391,6 +403,7 @@ def get_matching_codes(
             autofix_codes=autofix_codes,
             error_on_autofix=options.error_on_autofix,
             no_checkpoint_warning_decorators=options.no_checkpoint_warning_decorators,
+            transform_async_generator_decorators=options.transform_async_generator_decorators,
             exception_suppress_context_managers=options.exception_suppress_context_managers,
             startable_in_context_manager=options.startable_in_context_manager,
             async200_blocking_calls=options.async200_blocking_calls,

flake8_async/base.py

@@ -29,6 +29,7 @@ class Options:
     # whether to print an error message even when autofixed
     error_on_autofix: bool
     no_checkpoint_warning_decorators: Collection[str]
+    transform_async_generator_decorators: Collection[str]
     exception_suppress_context_managers: Collection[str]
     startable_in_context_manager: Collection[str]
     async200_blocking_calls: dict[str, str]

flake8_async/visitors/visitors.py

@@ -401,27 +401,32 @@ def leave_IfExp_test(self, node: cst.IfExp):
 @disabled_by_default
 class Visitor900(Flake8AsyncVisitor):
     error_codes: Mapping[str, str] = {
-        "ASYNC900": "Async generator without `@asynccontextmanager` not allowed."
+        "ASYNC900": "Async generator not allowed, unless transformed by a known decorator (one of: {})."
     }
 
     def __init__(self, *args: Any, **kwargs: Any):
         super().__init__(*args, **kwargs)
         self.unsafe_function: ast.AsyncFunctionDef | None = None
+        self.transform_decorators = (
+            "asynccontextmanager",
+            "fixture",
+            *self.options.transform_async_generator_decorators,
+        )
 
     def visit_AsyncFunctionDef(
         self, node: ast.AsyncFunctionDef | ast.FunctionDef | ast.Lambda
     ):
         self.save_state(node, "unsafe_function")
         if isinstance(node, ast.AsyncFunctionDef) and not has_decorator(
-            node, "asynccontextmanager", "fixture"
+            node, *self.transform_decorators
         ):
             self.unsafe_function = node
         else:
             self.unsafe_function = None
 
     def visit_Yield(self, node: ast.Yield):
         if self.unsafe_function is not None:
-            self.error(self.unsafe_function)
+            self.error(self.unsafe_function, ", ".join(self.transform_decorators))
             self.unsafe_function = None
 
     visit_FunctionDef = visit_AsyncFunctionDef

setup.py

@@ -3,7 +3,7 @@
 
 from pathlib import Path
 
-from setuptools import find_packages, setup
+from setuptools import find_packages, setup  # type: ignore
 
 
 def local_file(name: str) -> Path:

tests/eval_files/async900.py

@@ -54,3 +54,12 @@ async def cm():
 async def another_non_generator():
     def foo():
         yield
+
+
+# ARG --transform-async-generator-decorators=this_is_like_a_context_manager
+
+
+@this_is_like_a_context_manager()  # OK because of the config, issue #277
+async def some_generator():
+    while True:
+        yield