semanal: Check that async for/with is inside an async function (#12418)

* semanal: Disallow async for/with in non-async functions

* Fix invalid async syntax in test

Author: shoracek

mypy/message_registry.py

@@ -61,6 +61,9 @@ def format(self, *args: object, **kwargs: object) -> "ErrorMessage":
 )
 INCOMPATIBLE_TYPES_IN_ASYNC_FOR: Final = 'Incompatible types in "async for"'
 
+ASYNC_FOR_OUTSIDE_COROUTINE: Final = '"async for" outside async function'
+ASYNC_WITH_OUTSIDE_COROUTINE: Final = '"async with" outside async function'
+
 INCOMPATIBLE_TYPES_IN_YIELD: Final = ErrorMessage('Incompatible types in "yield"')
 INCOMPATIBLE_TYPES_IN_YIELD_FROM: Final = ErrorMessage('Incompatible types in "yield from"')
 INCOMPATIBLE_TYPES_IN_STR_INTERPOLATION: Final = "Incompatible types in string interpolation"

mypy/semanal.py

@@ -3621,6 +3621,10 @@ def visit_while_stmt(self, s: WhileStmt) -> None:
         self.visit_block_maybe(s.else_body)
 
     def visit_for_stmt(self, s: ForStmt) -> None:
+        if s.is_async:
+            if not self.is_func_scope() or not self.function_stack[-1].is_coroutine:
+                self.fail(message_registry.ASYNC_FOR_OUTSIDE_COROUTINE, s, code=codes.SYNTAX)
+
         self.statement = s
         s.expr.accept(self)
 
@@ -3680,6 +3684,10 @@ def visit_with_stmt(self, s: WithStmt) -> None:
         self.statement = s
         types: List[Type] = []
 
+        if s.is_async:
+            if not self.is_func_scope() or not self.function_stack[-1].is_coroutine:
+                self.fail(message_registry.ASYNC_WITH_OUTSIDE_COROUTINE, s, code=codes.SYNTAX)
+
         if s.unanalyzed_type:
             assert isinstance(s.unanalyzed_type, ProperType)
             actual_targets = [t for t in s.target if t is not None]
@@ -4175,12 +4183,24 @@ def visit_type_application(self, expr: TypeApplication) -> None:
                 expr.types[i] = analyzed
 
     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:
+                self.fail(message_registry.ASYNC_FOR_OUTSIDE_COROUTINE, expr, code=codes.SYNTAX)
+
         expr.generator.accept(self)
 
     def visit_set_comprehension(self, expr: SetComprehension) -> None:
+        if any(expr.generator.is_async):
+            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)
+
         expr.generator.accept(self)
 
     def visit_dictionary_comprehension(self, expr: DictionaryComprehension) -> None:
+        if any(expr.is_async):
+            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)
+
         with self.enter(expr):
             self.analyze_comp_for(expr)
             expr.key.accept(self)

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

@@ -825,3 +825,42 @@ def bar() -> None:
 
 [builtins fixtures/async_await.pyi]
 [typing fixtures/typing-async.pyi]
+
+[case testAsyncForOutsideCoroutine]
+# flags: --python-version 3.7
+
+async def g():
+    yield 0
+
+def f() -> None:
+    [x async for x in g()] # E: "async for" outside async function
+    {x async for x in g()} # E: "async for" outside async function
+    {x: True async for x in g()} # E: "async for" outside async function
+    (x async for x in g())
+    async for x in g(): ... # E: "async for" outside async function
+
+[x async for x in g()] # E: "async for" outside async function
+{x async for x in g()} # E: "async for" outside async function
+{x: True async for x in g()} # E: "async for" outside async function
+(x async for x in g())
+async for x in g(): ... # E: "async for" outside async function
+
+[builtins fixtures/async_await.pyi]
+[typing fixtures/typing-async.pyi]
+
+[case testAsyncWithOutsideCoroutine]
+# flags: --python-version 3.7
+
+class C:
+    async def __aenter__(self): pass
+    async def __aexit__(self, x, y, z): pass
+
+def f() -> None:
+    async with C() as x: # E: "async with" outside async function
+        pass
+
+async with C() as x: # E: "async with" outside async function
+    pass
+
+[builtins fixtures/async_await.pyi]
+[typing fixtures/typing-async.pyi]

test-data/unit/check-default-plugin.test

@@ -37,8 +37,9 @@ async def yield_id(item: T) -> AsyncGenerator[T, None]:
 
 reveal_type(yield_id) # N: Revealed type is "def [T] (item: T`-1) -> typing.AsyncContextManager[T`-1]"
 
-async with yield_id(1) as x:
-    reveal_type(x) # N: Revealed type is "builtins.int*"
+async def f() -> None:
+    async with yield_id(1) as x:
+        reveal_type(x) # N: Revealed type is "builtins.int*"
 [typing fixtures/typing-async.pyi]
 [builtins fixtures/tuple.pyi]