8000 Support type comments in async for/with (#2858) · python/mypy@4bc831f · GitHub
[go: up one dir, main page]

Skip to content

Commit 4bc831f

Browse files
afriederddfisher
authored andcommitted
Support type comments in async for/with (#2858)
* Add support for type comments in async for/with * Add tests for async for/with
1 parent 3cc5201 commit 4bc831f

File tree

2 files changed

+47
-5
lines changed

2 files changed

+47
-5
lines changed

mypy/fastparse.py

Lines changed: 14 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -532,13 +532,18 @@ def visit_For(self, n: ast3.For) -> ForStmt:
532532
self.as_block(n.orelse, n.lineno),
533533
target_type)
534534

535-
# AsyncFor(expr target, expr iter, stmt* body, stmt* orelse)
535+
# AsyncFor(expr target, expr iter, stmt* body, stmt* orelse, string? type_comment)
536536
@with_line
537537
def visit_AsyncFor(self, n: ast3.AsyncFor) -> ForStmt:
538+
if n.type_comment is not None:
539+
target_type = parse_type_comment(n.type_comment, n.lineno, self.errors)
540+
else:
541+
target_type = None
538542
r = ForStmt(self.visit(n.target),
539543
self.visit(n.iter),
540544
self.as_block(n.body, n.lineno),
541-
self.as_block(n.orelse, n.lineno))
545+
self.as_block(n.orelse, n.lineno),
546+
target_type)
542547
r.is_async = True
543548
return r
544549

@@ -568,12 +573,17 @@ def visit_With(self, n: ast3.With) -> WithStmt:
568573
self.as_block(n.body, n.lineno),
569574
target_type)
570575

571-
# AsyncWith(withitem* items, stmt* body)
576+
# AsyncWith(withitem* items, stmt* body, string? type_comment)
572577
@with_line
573578
def visit_AsyncWith(self, n: ast3.AsyncWith) -> WithStmt:
579+
if n.type_comment is not None:
580+
target_type = parse_type_comment(n.type_comment, n.lineno, self.errors)
581+
else:
582+
target_type = None
574583
r = WithStmt([self.visit(i.context_expr) for i in n.items],
575584
[self.visit(i.optional_vars) for i in n.items],
576-
self.as_block(n.body, n.lineno))
585+
self.as_block(n.body, n.lineno),
586+
target_type)
577587
r.is_async = True
578588
return r
579589

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

Lines changed: 33 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -152,6 +152,22 @@ async def f() -> None:
152152
main:4: error: AsyncIterable expected
153153
main:4: error: List[int] has no attribute "__aiter__"
154154

155+
[case testAsyncForTypeComments]
156+
# flags: --fast-parser
157+
from typing import AsyncIterator, Union
158+
class C(AsyncIterator[int]):
159+
async def __anext__(self) -> int: return 0
160+
async def f() -> None:
161+
async for x in C(): # type: str # E: Incompatible types in assignment (expression has type "int", variable has type "str")
162+
pass
163+
164+
async for y in C(): # type: int
165+
pass
166+
167+
async for z in C(): # type: Union[int, str]
168+
reveal_type(z) # E: Revealed type is 'Union[builtins.int, builtins.str]'
169+
[builtins fixtures/async_await.pyi]
170+
155171
[case testAsyncWith]
156172
# flags: --fast-parser
157173
class C:
@@ -161,7 +177,7 @@ async def f() -> None:
161177
async with C() as x:
162178
reveal_type(x) # E: Revealed type is 'builtins.int*'
163179
[builtins fixtures/async_await.pyi]
164-
[out]
180+
165181

166182
[case testAsyncWithError]
167183
# flags: --fast-parser
@@ -220,6 +236,22 @@ async def f() -> None:
220236
[builtins fixtures/async_await.pyi]
221237
[out]
222238

239+
[case testAsyncWithTypeComments]
240+
# flags: --fast-parser
241+
class C:
242+
async def __aenter__(self) -> int: pass
243+
async def __aexit__(self, x, y, z) -> None: pass
244+
async def f() -> None:
245+
async with C() as x: # type: int
246+
pass
247+
248+
async with C() as y, C() as z: # type: str, int # E: Incompatible types in assignment (expression has type "int", variable has type "str")
249+
pass
250+
251+
async with C() as a: # type: int, int # E: Invalid tuple literal type
252+
pass
253+
[builtins fixtures/async_await.pyi]
254+
223255
[case testNoYieldInAsyncDef]
224256
# flags: --fast-parser
225257
async def f():

0 commit comments

Comments
 (0)
0