Add config option for async900 decorators (#279)

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/rules.rst

@@ -161,12 +161,15 @@ 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. `@trio_util.trio_async_generator
+       <https://trio-util.readthedocs.io/en/latest/index.html#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,33 @@ 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

@@ -3,7 +3,7 @@
 from contextlib import asynccontextmanager
 
 
-async def foo1():  # ASYNC900: 0
+async def foo1():  # ASYNC900: 0, 'asynccontextmanager, fixture, this_is_like_a_context_manager'
     yield
     yield
 
@@ -15,7 +15,7 @@ async def foo2():
 
 @asynccontextmanager
 async def foo3():
-    async def bar():  # ASYNC900: 4
+    async def bar():  # ASYNC900: 4, 'asynccontextmanager, fixture, this_is_like_a_context_manager'
         yield
 
     yield
@@ -37,7 +37,7 @@ async def async_fixtures_can_take_arguments():
 
 # no-checkpoint-warning-decorator now ignored
 @other_context_manager
-async def foo5():  # ASYNC900: 0
+async def foo5():  # ASYNC900: 0, 'asynccontextmanager, fixture, this_is_like_a_context_manager'
     yield
 
 
@@ -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