add --exception-suppress-context-manager (#254)

* add --exception-suppress-context-manager

* add support for `from contextlib import suppress [as <xxx>]` and `from contextlib import *`

docs/usage.rst

@@ -56,6 +56,8 @@ install and run as standalone
 
 If inside a git repository, running without arguments will run it against all ``*.py`` files in the repository.
 
+Note that this does not currently support reading config files, and does not respect ``# noqa`` comments.
+
 .. code-block:: sh
 
    pip install flake8-async
@@ -230,6 +232,32 @@ Example
      ign*,
      *.ignore,
 
+``exception-suppress-context-managers``
+---------------------------------------
+
+
+Comma-separated list of contextmanagers which may suppress exceptions
+without reraising. For ASYNC91x, these will be parsed in the worst-case scenario,
+where any checkpoints inside the contextmanager are not executed, and all
+exceptions are suppressed.
+``contextlib.suppress`` will be added to the list after parsing, and some basic parsing
+of ``from contextlib import suppress`` is supported.
+Decorators can be dotted or not, as well as support * as a wildcard.
+
+If you want to be extremely pessimistic, you can specify ``*`` as the context manager.
+We may add a whitelist option in the future to support this use-case better.
+
+Example
+^^^^^^^
+
+.. code-block:: none
+
+   exception-suppress-context-managers =
+     mysuppressor,
+     dangerouslibrary.*,
+     *.suppress,
+
+
 .. _--startable-in-context-manager:
 
 ``startable-in-context-manager``

flake8_async/__init__.py

@@ -261,6 +261,18 @@ def add_options(option_manager: OptionManager | ArgumentParser):
                 "mydecorator,mypackage.mydecorators.*``"
             ),
         )
+        add_argument(
+            "--exception-suppress-context-managers",
+            default="",
+            required=False,
+            type=comma_separated_list,
+            help=(
+                "Comma-separated list of contextmanagers which may suppress exceptions "
+                "without reraising, breaking checkpoint guarantees of ASYNC91x. "
+                "``contextlib.suppress`` will be added to the list. "
+                "Decorators can be dotted or not, as well as support * as a wildcard. "
+            ),
+        )
         add_argument(
             "--startable-in-context-manager",
             type=parse_async114_identifiers,
@@ -379,6 +391,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,
+            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,
             anyio=options.anyio,

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]
+    exception_suppress_context_managers: Collection[str]
     startable_in_context_manager: Collection[str]
     async200_blocking_calls: dict[str, str]
     anyio: bool

flake8_async/visitors/helpers.py

@@ -112,7 +112,8 @@ def fnmatch_qualified_name(name_list: list[ast.expr], *patterns: str) -> str | N
 
 
 def fnmatch_qualified_name_cst(
-    name_list: Iterable[cst.Decorator], *patterns: str
+    name_list: Iterable[cst.Decorator | cst.Call | cst.Attribute | cst.Name],
+    *patterns: str,
 ) -> str | None:
     for name in name_list:
         qualified_name = get_full_name_for_node_or_raise(name)

flake8_async/visitors/visitor91x.py

@@ -283,6 +283,9 @@ def __init__(self, *args: Any, **kwargs: Any):
         self.has_checkpoint_stack: list[bool] = []
         self.node_dict: dict[cst.With, list[AttributeCall]] = {}
 
+        # --exception-suppress-context-manager
+        self.suppress_imported_as: list[str] = []
+
     def should_autofix(self, node: cst.CSTNode, code: str | None = None) -> bool:
         if code is None:  # pragma: no branch
             code = "ASYNC911" if self.has_yield else "ASYNC910"
@@ -300,6 +303,28 @@ def checkpoint(self) -> None:
     def checkpoint_statement(self) -> cst.SimpleStatementLine:
         return checkpoint_statement(self.library[0])
 
+    def visit_ImportFrom(self, node: cst.ImportFrom) -> None:
+        # Semi-crude approach to handle `from contextlib import suppress`.
+        # It does not handle the identifier being overridden, or assigned
+        # to other idefintifers. Function scoping is handled though.
+        # The "proper" way would be to add a cst version of
+        # visitor_utility.VisitorTypeTracker, and expand that to handle imports.
+        if isinstance(node.module, cst.Name) and node.module.value == "contextlib":
+            # handle `from contextlib import *`
+            if isinstance(node.names, cst.ImportStar):
+                self.suppress_imported_as.append("suppress")
+                return
+            for alias in node.names:
+                if alias.name.value == "suppress":
+                    if alias.asname is not None:
+                        # `libcst.AsName` is incorrectly typed
+                        # https://github.com/Instagram/LibCST/issues/503
+                        assert isinstance(alias.asname.name, cst.Name)
+                        self.suppress_imported_as.append(alias.asname.name.value)
+                    else:
+                        self.suppress_imported_as.append("suppress")
+                    return
+
     def visit_FunctionDef(self, node: cst.FunctionDef) -> bool:
         # don't lint functions whose bodies solely consist of pass or ellipsis
         # @overload functions are also guaranteed to be empty
@@ -316,6 +341,7 @@ def visit_FunctionDef(self, node: cst.FunctionDef) -> bool:
             "loop_state",
             "try_state",
             "has_checkpoint_stack",
+            "suppress_imported_as",
             copy=True,
         )
         self.uncheckpointed_statements = set()
@@ -449,12 +475,29 @@ def leave_Await(
     # can't use TypeVar due to libcst's built-in type checking not supporting it
     leave_Raise = leave_Await  # type: ignore
 
+    def _is_exception_suppressing_context_manager(self, node: cst.With) -> bool:
+        return (
+            fnmatch_qualified_name_cst(
+                (x.item for x in node.items if isinstance(x.item, cst.Call)),
+                "contextlib.suppress",
+                *self.suppress_imported_as,
+                *self.options.exception_suppress_context_managers,
+            )
+            is not None
+        )
+
     # Async context managers can reasonably checkpoint on either or both of entry and
     # exit.  Given that we can't tell which, we assume "both" to avoid raising a
     # missing-checkpoint warning when there might in fact be one (i.e. a false alarm).
     def visit_With_body(self, node: cst.With):
         if getattr(node, "asynchronous", None):
             self.checkpoint()
+
+        # if this might suppress exceptions, we cannot treat anything inside it as
+        # checkpointing.
+        if self._is_exception_suppressing_context_manager(node):
+            self.save_state(node, "uncheckpointed_statements", copy=True)
+
         if res := (
             with_has_call(node, *cancel_scope_names)
             or with_has_call(
@@ -499,6 +542,14 @@ def leave_With(self, original_node: cst.With, updated_node: cst.With):
                 self.uncheckpointed_statements.remove(s)
                 for res in self.node_dict[original_node]:
                     self.error(res.node, error_code="ASYNC912")
+
+        # if exception-suppressing, restore all uncheckpointed statements from
+        # before the `with`.
+        if self._is_exception_suppressing_context_manager(original_node):
+            prev_checkpoints = self.uncheckpointed_statements
+            self.restore_state(original_node)
+            self.uncheckpointed_statements.update(prev_checkpoints)
+
         if getattr(original_node, "asynchronous", None):
             self.checkpoint()
         return updated_node

tests/autofix_files/async910.py

@@ -3,6 +3,7 @@
 # mypy: disable-error-code="unreachable"
 from __future__ import annotations
 
+import contextlib
 import typing
 from typing import Any, overload
 

tests/autofix_files/exception_suppress_context_manager.py

@@ -0,0 +1,171 @@
+# ARG --exception-suppress-context-managers=mysuppress,*.dangerousname,dangerouslibrary.*
+# ARG --enable=ASYNC910,ASYNC911
+
+# AUTOFIX
+# ASYNCIO_NO_AUTOFIX
+
+# 912 is tested in eval_files/async912.py to avoid problems with autofix/asyncio
+
+import contextlib
+from typing import Any
+
+import trio
+
+mysuppress: Any
+anything: Any
+dangerouslibrary: Any
+
+
+async def foo() -> Any:
+    await foo()
+
+
+async def foo_suppress():  # ASYNC910: 0, "exit", Statement('function definition', lineno)
+    with contextlib.suppress():
+        await foo()
+    await trio.lowlevel.checkpoint()
+
+
+async def foo_suppress_1():  # ASYNC910: 0, "exit", Statement('function definition', lineno)
+    with mysuppress():
+        await foo()
+    await trio.lowlevel.checkpoint()
+
+
+async def foo_suppress_2():  # ASYNC910: 0, "exit", Statement('function definition', lineno)
+    with anything.dangerousname():
+        await foo()
+    await trio.lowlevel.checkpoint()
+
+
+async def foo_suppress_3():  # ASYNC910: 0, "exit", Statement('function definition', lineno)
+    with dangerouslibrary.anything():
+        await foo()
+    await trio.lowlevel.checkpoint()
+
+
+async def foo_suppress_async911():  # ASYNC911: 0, "exit", Statement("function definition", lineno)
+    with contextlib.suppress():
+        await foo()
+        yield
+        await foo()
+    await trio.lowlevel.checkpoint()
+
+
+# the `async with` checkpoints, so there's no error
+async def foo_suppress_async():
+    async with mysuppress:
+        await foo()
+
+
+async def foo_multiple():  # ASYNC910: 0, "exit", Statement('function definition', lineno)
+    with anything, contextlib.suppress():
+        await foo()
+    await trio.lowlevel.checkpoint()
+
+
+# we only match on *calls*
+async def foo_no_call():
+    with contextlib.suppress:  # type: ignore[attr-defined]
+        await foo()
+
+
+# doesn't work on namedexpr, but those should use `as` anyway
+async def foo_namedexpr():
+    with (ref := contextlib.suppress()):
+        await foo()
+
+
+async def foo_suppress_as():  # ASYNC910: 0, "exit", Statement('function definition', lineno)
+    with contextlib.suppress() as my_suppressor:
+        await foo()
+    await trio.lowlevel.checkpoint()
+
+
+# ###############################
+# from contextlib import suppress
+# ###############################
+
+
+# not enabled unless it's imported from contextlib
+async def foo_suppress_directly_imported_1():
+    with suppress():
+        await foo()
+
+
+from contextlib import suppress
+
+
+# now it's enabled
+async def foo_suppress_directly_imported_2():  # ASYNC910: 0, "exit", Statement('function definition', lineno)
+    with suppress():
+        await foo()
+    await trio.lowlevel.checkpoint()
+
+
+# it also supports importing with an alias
+from contextlib import suppress as adifferentname
+
+
+async def foo_suppress_directly_imported_3():  # ASYNC910: 0, "exit", Statement('function definition', lineno)
+    with adifferentname():
+        await foo()
+    await trio.lowlevel.checkpoint()
+
+
+# and will keep track of all identifiers it's been assigned as
+async def foo_suppress_directly_imported_4():  # ASYNC910: 0, "exit", Statement('function definition', lineno)
+    with suppress():
+        await foo()
+    await trio.lowlevel.checkpoint()
+
+
+# basic function scoping is supported
+async def function_that_contains_the_import():
+    from contextlib import suppress as bar
+
+    with adifferentname():
+        await foo()
+    await trio.lowlevel.checkpoint()
+    yield  # ASYNC911: 4, "yield", Stmt("function definition", lineno-5)
+    with bar():
+        await foo()
+    await trio.lowlevel.checkpoint()
+    yield  # ASYNC911: 4, "yield", Stmt("yield", lineno-3)
+    await foo()
+
+
+# bar is not suppressing
+async def foo_suppress_directly_imported_scoped():
+    with bar():  # type: ignore[name-defined]
+        await foo()
+
+
+# adifferentname is still suppressing
+async def foo_suppress_directly_imported_restored_after_scope():  # ASYNC910: 0, "exit", Statement('function definition', lineno)
+    with adifferentname():
+        await foo()
+    await trio.lowlevel.checkpoint()
+
+
+# We don't track the identifier being overridden though.
+adifferentname = None  # type: ignore[assignment]
+
+
+# shouldn't give an error
+async def foo_suppress_directly_imported_5():  # ASYNC910: 0, "exit", Statement('function definition', lineno)
+    with adifferentname():
+        await foo()
+    await trio.lowlevel.checkpoint()
+
+
+# or assignments to different identifiers
+from contextlib import suppress
+
+my_second_suppress = suppress
+
+
+# should give an error
+async def foo_suppress_directly_imported_assignment():
+    with my_second_suppress():
+        await foo()