fix: added await-outside-coroutine checks also to list and dict comprehensions

Author: josetapadas

mypy/message_registry.py

@@ -79,8 +79,8 @@ def with_additional_msg(self, info: str) -> ErrorMessage:
 
 ASYNC_FOR_OUTSIDE_COROUTINE: Final = '"async for" outside async function'
 ASYNC_WITH_OUTSIDE_COROUTINE: Final = '"async with" outside async function'
-
-AWAIT_WITH_OUTSIDE_COROUTINE: Final = '"await" outside coroutine ("async def")'
+AWAIT_OUTSIDE_FUNCTION: Final = '"await" outside function'
+AWAIT_OUTSIDE_COROUTINE: Final = '"await" outside coroutine ("async def")'
 
 INCOMPATIBLE_TYPES_IN_YIELD: Final = ErrorMessage('Incompatible types in "yield"')
 INCOMPATIBLE_TYPES_IN_YIELD_FROM: Final = ErrorMessage('Incompatible types in "yield from"')

mypy/semanal.py

@@ -6074,6 +6074,18 @@ def visit_type_application(self, expr: TypeApplication) -> None:
             if analyzed is not None:
                 expr.types[i] = analyzed
 
+    def _check_await_outside_coroutine(self, expr: ListComprehension | SetComprehension | DictionaryComprehension) -> None:
+        if not has_await_expression(expr):
+            return
+
+        if not self.is_func_scope() or not self.function_stack[-1].is_coroutine:
+            self.fail(
+                message_registry.AWAIT_OUTSIDE_COROUTINE,
+                expr,
+                code=codes.AWAIT_NOT_ASYNC,
+                serious=True,
+            )
+
     def visit_list_comprehension(self, expr: ListComprehension) -> None:
         if any(expr.generator.is_async):
             if not self.is_func_scope() or not self.function_stack[-1].is_coroutine:
@@ -6083,14 +6095,8 @@ def visit_list_comprehension(self, expr: ListComprehension) -> None:
                     code=codes.SYNTAX,
                     serious=True,
                 )
-        elif has_await_expression(expr):
-            if not self.is_func_scope() or not self.function_stack[-1].is_coroutine:
-                self.fail(
-                    message_registry.AWAIT_WITH_OUTSIDE_COROUTINE,
-                    expr,
-                    code=codes.AWAIT_NOT_ASYNC,
-                    serious=True,
-                )
+
+        self._check_await_outside_coroutine(expr)
 
         expr.generator.accept(self)
 
@@ -6099,6 +6105,8 @@ def visit_set_comprehension(self, expr: SetComprehension) -> None:
             if not self.is_func_scope() or not self.function_stack[-1].is_coroutine:
                 self.fail(message_registry.ASYNC_FOR_OUTSIDE_COROUTINE, expr, code=codes.SYNTAX)
 
+        self._check_await_outside_coroutine(expr)
+
         expr.generator.accept(self)
 
     def visit_dictionary_comprehension(self, expr: DictionaryComprehension) -> None:
@@ -6111,6 +6119,8 @@ def visit_dictionary_comprehension(self, expr: DictionaryComprehension) -> None:
                     serious=True,
                 )
 
+        self._check_await_outside_coroutine(expr)
+
         with self.enter(expr):
             self.analyze_comp_for(expr)
             expr.key.accept(self)
@@ -6185,13 +6195,10 @@ def visit_await_expr(self, expr: AwaitExpr) -> None:
             # We check both because is_function_scope() returns True inside comprehensions.
             # This is not a blocker, because some enviroments (like ipython)
             # support top level awaits.
-            self.fail('"await" outside function', expr, serious=True, code=codes.TOP_LEVEL_AWAIT)
-        elif (
-            not self.function_stack[-1].is_coroutine
-            and self.scope_stack[-1] != SCOPE_COMPREHENSION
-        ):
+            self.fail(message_registry.AWAIT_OUTSIDE_FUNCTION, expr, serious=True, code=codes.TOP_LEVEL_AWAIT)
+        elif not self.function_stack[-1].is_coroutine and self.scope_stack[-1] != SCOPE_COMPREHENSION:
             self.fail(
-                message_registry.AWAIT_WITH_OUTSIDE_COROUTINE,
+                message_registry.AWAIT_OUTSIDE_COROUTINE,
                 expr,
                 serious=True,
                 code=codes.AWAIT_NOT_ASYNC,

test-data/unit/check-async-await.test

@@ -1085,12 +1085,23 @@ class Launcher(P):
 [builtins fixtures/async_await.pyi]
 [typing fixtures/typing-async.pyi]
 
-[case testAwaitInsideGeneratorExpr]
+[case testAwaitOutsideCoroutine]
 def foo():
     yield 0
 
 def bar():
     (await x for x in foo())
 
+    (await x for x in xs)  # OK
+    [await x for x in xs]  # E: "await" outside coroutine ("async def")
+    {await x for x in xs}  # E: "await" outside coroutine ("async def")
+    {0: await x for x in xs} # E: "await" outside coroutine ("async def")
+    {await x: 0 for x in xs} # E: "await" outside coroutine ("async def")
+
+    (x for x in await xs) # E: "await" outside coroutine ("async def")
+    [x for x in await xs] # E: "await" outside coroutine ("async def")
+    {x for x in await xs} # E: "await" outside coroutine ("async def")
+    {x: 0 for x in await xs} # E: "await" outside coroutine ("async def")
+
 [builtins fixtures/async_await.pyi]
 [typing fixtures/typing-async.pyi]