When copying a ForStmt or WithStmt, copy is_async flag (python#9268)

gvanrossum · web-flow · commit 781caffa5c32 · 2020-08-05T07:59:30.000-07:00
(The test only tests 'async with', but 'async for' had the same bug.) Fixes python#9261.
diff --git a/mypy/treetransform.py b/mypy/treetransform.py
@@ -250,6 +250,7 @@ def visit_for_stmt(self, node: ForStmt) -> ForStmt:
                       self.block(node.body),
                       self.optional_block(node.else_body),
                       self.optional_type(node.unanalyzed_index_type))
+        new.is_async = node.is_async
         new.index_type = self.optional_type(node.index_type)
         return new
 
@@ -293,6 +294,7 @@ def visit_with_stmt(self, node: WithStmt) -> WithStmt:
                        self.optional_expressions(node.target),
                        self.block(node.body),
                        self.optional_type(node.unanalyzed_type))
+        new.is_async = node.is_async
         new.analyzed_types = [self.type(typ) for typ in node.analyzed_types]
         return new
 
diff --git a/test-data/unit/check-async-await.test b/test-data/unit/check-async-await.test
@@ -745,3 +745,16 @@ def t() -> None:
 
 [builtins fixtures/async_await.pyi]
 [typing fixtures/typing-async.pyi]
+
+[case testAsyncWithInGenericClass]
+from typing import Generic, AsyncContextManager, TypeVar
+
+T = TypeVar('T', str, int)
+
+class Foo(Generic[T]):
+    async def foo(self, manager: AsyncContextManager):
+        async with manager:
+            pass
+
+[builtins fixtures/async_await.pyi]
+[typing fixtures/typing-async.pyi]
diff --git a/test-data/unit/fixtures/typing-async.pyi b/test-data/unit/fixtures/typing-async.pyi
@@ -118,3 +118,8 @@ class ContextManager(Generic[T]):
     def __enter__(self) -> T: pass
     # Use Any because not all the precise types are in the fixtures.
     def __exit__(self, exc_type: Any, exc_value: Any, traceback: Any) -> Any: pass
+
+class AsyncContextManager(Generic[T]):
+    def __aenter__(self) -> Awaitable[T]: pass
+    # Use Any because not all the precise types are in the fixtures.
+    def __aexit__(self, exc_type: Any, exc_value: Any, traceback: Any) -> Awaitable[Any]: pass
