add comprehension support to 91x

Author: jakkdl

CHANGELOG.md

@@ -1,6 +1,9 @@
 # Changelog
 *[CalVer, YY.month.patch](https://calver.org/)*
 
+## Future
+- TRIO91X now supports comprehensions
+
 ## 23.2.5
 - Fix false alarms for `@pytest.fixture`-decorated functions in TRIO101, TRIO910 and TRIO911
 

flake8_trio/visitors/visitor91x.py

@@ -97,6 +97,7 @@ def __init__(self, *args: Any, **kwargs: Any):
         self.safe_decorator = False
         self.async_function = False
         self.uncheckpointed_statements: set[Statement] = set()
+        self.comp_unknown = False
 
         self.loop_state = LoopState()
         self.try_state = TryState()
@@ -522,3 +523,70 @@ def leave_BooleanOperation_right(self, node: cst.BooleanOperation):
         self.uncheckpointed_statements.update(
             self.outer[node]["uncheckpointed_statements"]
         )
+
+    # comprehensions are simpler than loops, since they cannot contain yields
+    # or many other complicated statements, but their subfields are not in the order
+    # they're logically executed, so we manually visit each field in execution order,
+    # as long as the effect of the statement is not known. Once we know the comprehension
+    # will checkpoint, we stop visiting, or once we are no longer guaranteed to execute
+    # code deeper in the comprehension.
+    # Functions return `False` so libcst doesn't iterate subnodes [again].
+    def visit_ListComp(self, node: cst.DictComp | cst.SetComp | cst.ListComp):
+        if not self.async_function or not self.uncheckpointed_statements:
+            return False
+
+        outer = self.comp_unknown
+        self.comp_unknown = True
+
+        # visit `for` and `if`s
+        node.for_in.visit(self)
+
+        # if still unknown, visit the expression
+        if self.comp_unknown and self.uncheckpointed_statements:
+            if isinstance(node, cst.DictComp):
+                node.key.visit(self)
+                node.value.visit(self)
+            else:
+                node.elt.visit(self)
+
+        self.comp_unknown = outer
+        return False
+
+    visit_SetComp = visit_ListComp
+    visit_DictComp = visit_ListComp
+
+    def visit_CompFor(self, node: cst.CompFor):
+        # should only ever be visited manually, when inside an async function
+        assert self.async_function
+
+        if not self.uncheckpointed_statements:
+            return False
+
+        # if async comprehension, checkpoint
+        if node.asynchronous:
+            self.uncheckpointed_statements = set()
+            self.comp_unknown = False
+            return False
+
+        # visit the iter call, which might have await's
+        node.iter.visit(self)
+
+        # stop checking if the loop is not guaranteed to execute
+        if not iter_guaranteed_once_cst(node.iter):
+            self.comp_unknown = False
+
+        # only the first `if` is guaranteed to execute
+        # and if there's any, don't check inner loop
+        elif node.ifs:
+            self.comp_unknown = False
+            node.ifs[0].visit(self)
+        elif node.inner_for_in:
+            # visit nested loops (and ifs), if any
+            node.inner_for_in.visit(self)
+
+        return False
+
+    # We don't have any logic on if generators are guaranteed to unroll, so always
+    # ignore their content
+    def visit_GeneratorExp(self, node: cst.GeneratorExp):
+        return False

tests/eval_files/trio910.py

@@ -543,10 +543,9 @@ async def foo_comprehension_1():
 
 
 # should error
-async def foo_comprehension_2():
+async def foo_comprehension_2():  # error: 0, "exit", Statement("function definition", lineno)
     [await foo() for x in range(10) if bar()]
 
 
-# should not error, see https://github.com/Zac-HD/flake8-trio/issues/144
-async def foo_comprehension_3():  # error: 0, "exit", Statement("function definition", lineno)
+async def foo_comprehension_3():
     [... async for x in bar()]

tests/eval_files/trio911.py

@@ -12,6 +12,10 @@ async def foo() -> Any:
     await foo()
 
 
+def bar(*args) -> Any:
+    ...
+
+
 async def foo_yield_1():
     await foo()
     yield 5
@@ -886,3 +890,97 @@ async def foo_test():
 @pytest.fixture()
 async def foo_test2():
     yield
+
+
+async def comprehensions():
+    # guaranteed iteration with await in test
+    [... for x in range(10) if await foo()]
+    yield  # safe
+
+    # guaranteed iteration and await in value, but test is not guaranteed
+    [await foo() for x in range(10) if bar()]
+    yield  # error: 4, "yield", Stmt("yield", line-4)
+
+    # guaranteed iteration and await in value
+    [await foo() for x in range(10)]
+    yield  # safe
+
+    # not guaranteed to iter
+    [await foo() for x in bar()]
+    yield  # error: 4, "yield", Stmt("yield", line-4)
+
+    # await statement in loop expression
+    [... for x in bar(await foo())]
+    yield
+
+    # set comprehensions use same logic as list
+    {await foo() for x in range(10)}
+    yield  # safe
+
+    {await foo() for x in bar()}
+    yield  # error: 4, "yield", Stmt("yield", line-3)
+
+    # dict comprehensions use same logic as list
+    {await foo(): 5 for x in bar()}
+    yield  # error: 4, "yield", Stmt("yield", line-4)
+
+    # other than `await` can be in both key&val
+    {await foo(): 5 for x in range(10)}
+    yield
+
+    {5: await foo() for x in range(10)}
+    yield
+
+    # generator expressions are never treated as safe
+    (await foo() for x in range(10))
+    yield  # error: 4, "yield", Stmt("yield", line-4)
+
+    (await foo() for x in bar() if await foo())
+    yield  # error: 4, "yield", Stmt("yield", line-3)
+
+    # async for always safe
+    [... async for x in bar()]
+    yield  # safe
+    {... async for x in bar()}
+    yield  # safe
+    {...: ... async for x in bar()}
+    yield  # safe
+
+    # other than in generator expression
+    (... async for x in bar())
+    yield  # error: 4, "yield", Stmt("yield", line-4)
+
+    # multiple loops
+    [... for x in range(10) for y in range(10) if await foo()]
+    yield
+    [... for x in range(10) for y in bar() if await foo()]
+    yield  # error: 4, "yield", Stmt("yield", line-2)
+    [... for x in bar() for y in range(10) if await foo()]
+    yield  # error: 4, "yield", Stmt("yield", line-2)
+
+    [await foo() for x in range(10) for y in range(10)]
+    yield
+    [await foo() for x in range(10) for y in bar()]
+    yield  # error: 4, "yield", Stmt("yield", line-2)
+    [await foo() for x in bar() for y in range(10)]
+    yield  # error: 4, "yield", Stmt("yield", line-2)
+
+    # trip loops!
+    [... for x in range(10) for y in range(10) async for z in bar()]
+    yield
+    [... for x in range(10) for y in range(10) for z in range(10)]
+    yield  # error: 4, "yield", Stmt("yield", line-2)
+
+    # multiple ifs
+    [... for x in range(10) for y in range(10) if await foo() if await foo()]
+    yield
+
+    [... for x in range(10) for y in bar() if await foo() if await foo()]
+    yield  # error: 4, "yield", Stmt("yield", line-3)
+
+    # nested comprehensions
+    [[await foo() for x in range(10)] for y in range(10)]
+    yield
+
+    # code coverage: inner comprehension with no checkpointed statements
+    [... for x in [await foo()] for y in x]