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. :func:`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,19 +401,26 @@ 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"
     }
 
     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
+        )
+        self.error_codes["ASYNC900"] = (  # type: ignore
+            "Async generator not allowed, unless transformed by a known decorator"
+            f" (one of: {', '.join(self.transform_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:

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