From 0b6586c2e097a6f435886b18cc915a86074847e7 Mon Sep 17 00:00:00 2001
From: Alex Frieder <alex.frieder@gmail.com>
Date: Sun, 12 Feb 2017 19:47:24 -0500
Subject: [PATCH 1/2] Add support for type comments in async for/with

---
 mypy/fastparse.py | 18 ++++++++++++++----
 1 file changed, 14 insertions(+), 4 deletions(-)

diff --git a/mypy/fastparse.py b/mypy/fastparse.py
index b70fdd814ce8..9a4562298270 100644
--- a/mypy/fastparse.py
+++ b/mypy/fastparse.py
@@ -532,13 +532,18 @@ def visit_For(self, n: ast3.For) -> ForStmt:
                        self.as_block(n.orelse, n.lineno),
                        target_type)
 
-    # AsyncFor(expr target, expr iter, stmt* body, stmt* orelse)
+    # AsyncFor(expr target, expr iter, stmt* body, stmt* orelse, string? type_comment)
     @with_line
     def visit_AsyncFor(self, n: ast3.AsyncFor) -> ForStmt:
+        if n.type_comment is not None:
+            target_type = parse_type_comment(n.type_comment, n.lineno, self.errors)
+        else:
+            target_type = None
         r = ForStmt(self.visit(n.target),
                     self.visit(n.iter),
                     self.as_block(n.body, n.lineno),
-                    self.as_block(n.orelse, n.lineno))
+                    self.as_block(n.orelse, n.lineno),
+                    target_type)
         r.is_async = True
         return r
 
@@ -568,12 +573,17 @@ def visit_With(self, n: ast3.With) -> WithStmt:
                         self.as_block(n.body, n.lineno),
                         target_type)
 
-    # AsyncWith(withitem* items, stmt* body)
+    # AsyncWith(withitem* items, stmt* body, string? type_comment)
     @with_line
     def visit_AsyncWith(self, n: ast3.AsyncWith) -> WithStmt:
+        if n.type_comment is not None:
+            target_type = parse_type_comment(n.type_comment, n.lineno, self.errors)
+        else:
+            target_type = None
         r = WithStmt([self.visit(i.context_expr) for i in n.items],
                      [self.visit(i.optional_vars) for i in n.items],
-                     self.as_block(n.body, n.lineno))
+                     self.as_block(n.body, n.lineno),
+                     target_type)
         r.is_async = True
         return r
 

From 907a3b739ce153051efbc0b2dfedda6c4f189af3 Mon Sep 17 00:00:00 2001
From: Alex Frieder <alex.frieder@gmail.com>
Date: Sun, 12 Feb 2017 19:52:45 -0500
Subject: [PATCH 2/2] Add tests for async for/with

---
 test-data/unit/check-async-await.test | 34 ++++++++++++++++++++++++++-
 1 file changed, 33 insertions(+), 1 deletion(-)

diff --git a/test-data/unit/check-async-await.test b/test-data/unit/check-async-await.test
index e2a18bc0555b..34c983166e25 100644
--- a/test-data/unit/check-async-await.test
+++ b/test-data/unit/check-async-await.test
@@ -152,6 +152,22 @@ async def f() -> None:
 main:4: error: AsyncIterable expected
 main:4: error: List[int] has no attribute "__aiter__"
 
+[case testAsyncForTypeComments]
+# flags: --fast-parser
+from typing import AsyncIterator, Union
+class C(AsyncIterator[int]):
+    async def __anext__(self) -> int: return 0
+async def f() -> None:
+    async for x in C():  # type: str  # E: Incompatible types in assignment (expression has type "int", variable has type "str")
+        pass
+
+    async for y in C():  # type: int
+        pass
+
+    async for z in C():  # type: Union[int, str]
+        reveal_type(z)  # E: Revealed type is 'Union[builtins.int, builtins.str]'
+[builtins fixtures/async_await.pyi]
+
 [case testAsyncWith]
 # flags: --fast-parser
 class C:
@@ -161,7 +177,7 @@ async def f() -> None:
     async with C() as x:
         reveal_type(x)  # E: Revealed type is 'builtins.int*'
 [builtins fixtures/async_await.pyi]
-[out]
+
 
 [case testAsyncWithError]
 # flags: --fast-parser
@@ -220,6 +236,22 @@ async def f() -> None:
 [builtins fixtures/async_await.pyi]
 [out]
 
+[case testAsyncWithTypeComments]
+# flags: --fast-parser
+class C:
+    async def __aenter__(self) -> int: pass
+    async def __aexit__(self, x, y, z) -> None: pass
+async def f() -> None:
+    async with C() as x:  # type: int
+        pass
+
+    async with C() as y, C() as z:  # type: str, int  # E: Incompatible types in assignment (expression has type "int", variable has type "str")
+        pass
+
+    async with C() as a:  # type: int, int  # E: Invalid tuple literal type
+        pass
+[builtins fixtures/async_await.pyi]
+
 [case testNoYieldInAsyncDef]
 # flags: --fast-parser
 async def f():