python-trio · jakkdl · Oct 28, 2024 · Oct 25, 2024 · Oct 28, 2024
diff --git a/docs/changelog.rst b/docs/changelog.rst
@@ -4,6 +4,10 @@ Changelog
 
 `CalVer, YY.month.patch <https://calver.org/>`_
 
+24.10.2
+=======
+- :ref:`ASYNC102 <async102>` now also warns about ``await()`` inside ``__aexit__``.
+
 24.10.1
 =======
 - Add :ref:`ASYNC123 <async123>` bad-exception-group-flattening

diff --git a/docs/rules.rst b/docs/rules.rst
@@ -21,8 +21,8 @@ ASYNC101 : yield-in-cancel-scope
     This has substantial overlap with :ref:`ASYNC119 <ASYNC119>`, which will warn on almost all instances of ASYNC101, but ASYNC101 is about a conceptually different problem that will not get resolved by `PEP 533 <https://peps.python.org/pep-0533/>`_.
 
 _`ASYNC102` : await-in-finally-or-cancelled
-    ``await`` inside ``finally`` or :ref:`cancelled-catching <cancelled>` ``except:`` must have shielded :ref:`cancel scope <cancel_scope>` with timeout.
-    If not, the async call will immediately raise a new cancellation, suppressing the cancellation that was caught.
+    ``await`` inside ``finally``, :ref:`cancelled-catching <cancelled>` ``except:``, or ``__aexit__`` must have shielded :ref:`cancel scope <cancel_scope>` with timeout.
+    If not, the async call will immediately raise a new cancellation, suppressing any cancellation that was caught.
     See :ref:`ASYNC120 <async120>` for the general case where other exceptions might get suppressed.
     This is currently not able to detect asyncio shields.
 

diff --git a/flake8_async/__init__.py b/flake8_async/__init__.py
@@ -38,7 +38,7 @@
 
 
 # CalVer: YY.month.patch, e.g. first release of July 2022 == "22.7.1"
-__version__ = "24.10.1"
+__version__ = "24.10.2"
 
 
 # taken from https://github.com/Zac-HD/shed

diff --git a/flake8_async/visitors/__init__.py b/flake8_async/visitors/__init__.py
@@ -31,7 +31,7 @@
     visitor2xx,
     visitor91x,
     visitor101,
-    visitor102,
+    visitor102_120,
     visitor103_104,
     visitor105,
     visitor111,

diff --git a/flake8_async/visitors/visitor102.py → flake8_async/visitors/visitor102_120.py b/flake8_async/visitors/visitor102.py → flake8_async/visitors/visitor102_120.py
@@ -1,4 +1,8 @@
-"""Visitor102, which warns on unprotected `await` inside `finally`.
+"""Contains Visitor102 with ASYNC102 and ASYNC120.
+
+ASYNC102: await-in-finally-or-cancelled
+ASYNC120: await-in-except
+
 
 To properly protect they must be inside a shielded cancel scope with a timeout.
 """
@@ -222,6 +226,10 @@ def visit_FunctionDef(
 
         self._potential_120 = []
 
+        # lambda doesn't have `name` attribute
+        if getattr(node, "name", None) == "__aexit__":
+            self._critical_scope = Statement("__aexit__", node.lineno, node.col_offset)
+
     visit_AsyncFunctionDef = visit_FunctionDef
     # lambda can't contain await, try, except, raise, with, or assignments.
     # You also can't do assignment expressions with attributes. So we don't need to

diff --git a/pyproject.toml b/pyproject.toml
@@ -65,6 +65,7 @@ ignore = [
   "D105",
   "D106",
   "D107",
+  "D400",  # ends-in-period, stricter version of ends-in-punctuation
   "S101",
   "D203",  # one-blank-line-before-class
   "D213",  # multi-line-summary-second-line

diff --git a/tests/eval_files/async102.py b/tests/eval_files/async102.py
@@ -304,3 +304,8 @@ async def foo_nested_cs():
             await foo()  # error: 12, Statement("bare except", lineno-12)
             cs1.shield = True
             await foo()
+
+
+# treat __aexit__ as a critical scope
+async def __aexit__():
+    await foo()  # error: 4, Statement("__aexit__", lineno-1)