Implement 'async def' and friends ('await', 'async for', 'async with') (#1808)

* PEP 492 syntax: `async def` and `await`.

* Fix type errors; add async for and async with (not fully fledged).

* Dispose of Async{For,With}Stmt -- use is_async flag instead.

* Basic `async for` is working.

* Clear unneeded TODOs.

* Fledgeling `async with` support.

* Disallow `yield [from]` in `async def`.

* Check Python version before looking up typing.Awaitable.

Ensure 'async def' is a syntax error in Python 2 (at least with the fast parser).

* Vast strides in accuracy for visit_await_expr().

* Add `@with_line` to PEP 492 visit function definitions.

* Fix tests now that errors have line numbers.

* Tweak tests for async/await a bit.

* Get rid of remaining XXX issues.

is_generator_return_type() now takes an extra is_coroutine flag.

* Move PEP 492 nodes back where they belong.

* Respond to code review.

* Add tests expecting errors from async for/with.

* Test that await <generator> is an error.

* Verify that `yield from` does not accept coroutines.

This revealed a spurious error "Function does not return a value", fixed that.

* Disallow return value in generator declared as -> Iterator.

* Fix typo in comment.

* Refactor visit_with_stmt() into separate helper methods for async and regular.

Also Use get_generator_return_type() instead of manually unpacking the value.

* Fix lint error. Correct comment about default ts/tr.

* Improve errors when __aenter__/__aexit__ are not async. With tests.

* Refactor: move all extraction of T from Awaitable[T] to a single helper.

* Follow __await__ to extract t from subclass of Awaitable[t].

* Make get_generator_return_type() default to AnyType() (i.e. as it was).

* Fix test to match reverting get_generator_return_type() to default to Any.

* Rename get_awaitable_return_type() to check_awaitable_expr(), update docstring.

Author: gvanrossum

mypy/checker.py

@@ -14,9 +14,12 @@
     UnaryExpr, FuncExpr, ComparisonExpr,
     StarExpr, YieldFromExpr, NonlocalDecl, DictionaryComprehension,
     SetComprehension, ComplexExpr, EllipsisExpr, YieldExpr, Argument,
+    AwaitExpr,
     ARG_POS, ARG_OPT, ARG_STAR, ARG_NAMED, ARG_STAR2
 )
-from mypy.types import Type, CallableType, AnyType, UnboundType, TupleType, TypeList, EllipsisType
+from mypy.types import (
+    Type, CallableType, FunctionLike, AnyType, UnboundType, TupleType, TypeList, EllipsisType,
+)
 from mypy import defaults
 from mypy import experiments
 from mypy.errors import Errors
@@ -242,6 +245,17 @@ def visit_Module(self, mod: ast35.Module) -> Node:
     #              arg? kwarg, expr* defaults)
     @with_line
     def visit_FunctionDef(self, n: ast35.FunctionDef) -> Node:
+        return self.do_func_def(n)
+
+    # AsyncFunctionDef(identifier name, arguments args,
+    #                  stmt* body, expr* decorator_list, expr? returns, string? type_comment)
+    @with_line
+    def visit_AsyncFunctionDef(self, n: ast35.AsyncFunctionDef) -> Node:
+        return self.do_func_def(n, is_coroutine=True)
+
+    def do_func_def(self, n: Union[ast35.FunctionDef, ast35.AsyncFunctionDef],
+                    is_coroutine: bool = False) -> Node:
+        """Helper shared between visit_FunctionDef and visit_AsyncFunctionDef."""
         args = self.transform_args(n.args, n.lineno)
 
         arg_kinds = [arg.kind for arg in args]
@@ -285,6 +299,9 @@ def visit_FunctionDef(self, n: ast35.FunctionDef) -> Node:
                        args,
                        self.as_block(n.body, n.lineno),
                        func_type)
+        if is_coroutine:
+            # A coroutine is also a generator, mostly for internal reasons.
+            func_def.is_generator = func_def.is_coroutine = True
         if func_type is not None:
             func_type.definition = func_def
             func_type.line = n.lineno
@@ -345,9 +362,6 @@ def make_argument(arg: ast35.arg, default: Optional[ast35.expr], kind: int) -> A
 
         return new_args
 
-    # TODO: AsyncFunctionDef(identifier name, arguments args,
-    #                  stmt* body, expr* decorator_list, expr? returns, string? type_comment)
-
     def stringify_name(self, n: ast35.AST) -> str:
         if isinstance(n, ast35.Name):
             return n.id
@@ -419,7 +433,16 @@ def visit_For(self, n: ast35.For) -> Node:
                        self.as_block(n.body, n.lineno),
                        self.as_block(n.orelse, n.lineno))
 
-    # TODO: AsyncFor(expr target, expr iter, stmt* body, stmt* orelse)
+    # AsyncFor(expr target, expr iter, stmt* body, stmt* orelse)
+    @with_line
+    def visit_AsyncFor(self, n: ast35.AsyncFor) -> Node:
+        r = ForStmt(self.visit(n.target),
+                    self.visit(n.iter),
+                    self.as_block(n.body, n.lineno),
+                    self.as_block(n.orelse, n.lineno))
+        r.is_async = True
+        return r
+
     # While(expr test, stmt* body, stmt* orelse)
     @with_line
     def visit_While(self, n: ast35.While) -> Node:
@@ -441,7 +464,14 @@ def visit_With(self, n: ast35.With) -> Node:
                         [self.visit(i.optional_vars) for i in n.items],
                         self.as_block(n.body, n.lineno))
 
-    # TODO: AsyncWith(withitem* items, stmt* body)
+    # AsyncWith(withitem* items, stmt* body)
+    @with_line
+    def visit_AsyncWith(self, n: ast35.AsyncWith) -> Node:
+        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))
+        r.is_async = True
+        return r
 
     # Raise(expr? exc, expr? cause)
     @with_line
@@ -628,7 +658,11 @@ def visit_GeneratorExp(self, n: ast35.GeneratorExp) -> GeneratorExpr:
                              iters,
                              ifs_list)
 
-    # TODO: Await(expr value)
+    # Await(expr value)
+    @with_line
+    def visit_Await(self, n: ast35.Await) -> Node:
+        v = self.visit(n.value)
+        return AwaitExpr(v)
 
     # Yield(expr? value)
     @with_line

mypy/fastparse.py

@@ -40,6 +40,11 @@
 INCOMPATIBLE_TYPES = 'Incompatible types'
 INCOMPATIBLE_TYPES_IN_ASSIGNMENT = 'Incompatible types in assignment'
 INCOMPATIBLE_REDEFINITION = 'Incompatible redefinition'
+INCOMPATIBLE_TYPES_IN_AWAIT = 'Incompatible types in await'
+INCOMPATIBLE_TYPES_IN_ASYNC_WITH_AENTER = 'Incompatible types in "async with" for __aenter__'
+INCOMPATIBLE_TYPES_IN_ASYNC_WITH_AEXIT = 'Incompatible types in "async with" for __aexit__'
+INCOMPATIBLE_TYPES_IN_ASYNC_FOR = 'Incompatible types in "async for"'
+
 INCOMPATIBLE_TYPES_IN_YIELD = 'Incompatible types in yield'
 INCOMPATIBLE_TYPES_IN_YIELD_FROM = 'Incompatible types in "yield from"'
 INCOMPATIBLE_TYPES_IN_STR_INTERPOLATION = 'Incompatible types in string interpolation'
@@ -57,6 +62,7 @@
 INCOMPATIBLE_VALUE_TYPE = 'Incompatible dictionary value type'
 NEED_ANNOTATION_FOR_VAR = 'Need type annotation for variable'
 ITERABLE_EXPECTED = 'Iterable expected'
+ASYNC_ITERABLE_EXPECTED = 'AsyncIterable expected'
 INCOMPATIBLE_TYPES_IN_FOR = 'Incompatible types in for statement'
 INCOMPATIBLE_ARRAY_VAR_ARGS = 'Incompatible variable arguments in call'
 INVALID_SLICE_INDEX = 'Slice index must be an integer or None'

mypy/messages.py

@@ -416,6 +416,7 @@ class FuncItem(FuncBase):
     # Is this an overload variant of function with more than one overload variant?
     is_overload = False
     is_generator = False   # Contains a yield statement?
+    is_coroutine = False   # Defined using 'async def' syntax?
     is_static = False      # Uses @staticmethod?
     is_class = False       # Uses @classmethod?
     # Variants of function with type variables with values expanded
@@ -486,6 +487,7 @@ def serialize(self) -> JsonDict:
                 'is_property': self.is_property,
                 'is_overload': self.is_overload,
                 'is_generator': self.is_generator,
+                'is_coroutine': self.is_coroutine,
                 'is_static': self.is_static,
                 'is_class': self.is_class,
                 'is_decorated': self.is_decorated,
@@ -507,6 +509,7 @@ def deserialize(cls, data: JsonDict) -> 'FuncDef':
         ret.is_property = data['is_property']
         ret.is_overload = data['is_overload']
         ret.is_generator = data['is_generator']
+        ret.is_coroutine = data['is_coroutine']
         ret.is_static = data['is_static']
         ret.is_class = data['is_class']
         ret.is_decorated = data['is_decorated']
@@ -798,6 +801,7 @@ class ForStmt(Statement):
     expr = None  # type: Expression
     body = None  # type: Block
     else_body = None  # type: Block
+    is_async = False  # True if `async for ...` (PEP 492, Python 3.5)
 
     def __init__(self, index: Expression, expr: Expression, body: Block,
                  else_body: Block) -> None:
@@ -908,6 +912,7 @@ class WithStmt(Statement):
     expr = None  # type: List[Expression]
     target = None  # type: List[Expression]
     body = None  # type: Block
+    is_async = False  # True if `async with ...` (PEP 492, Python 3.5)
 
     def __init__(self, expr: List[Expression], target: List[Expression],
                  body: Block) -> None:
@@ -1705,6 +1710,18 @@ def accept(self, visitor: NodeVisitor[T]) -> T:
         return visitor.visit__promote_expr(self)
 
 
+class AwaitExpr(Node):
+    """Await expression (await ...)."""
+
+    expr = None  # type: Node
+
+    def __init__(self, expr: Node) -> None:
+        self.expr = expr
+
+    def accept(self, visitor: NodeVisitor[T]) -> T:
+        return visitor.visit_await_expr(self)
+
+
 # Constants
 
 

mypy/nodes.py

@@ -957,6 +957,10 @@ def parse_statement(self) -> Tuple[Node, bool]:
             stmt = self.parse_exec_stmt()
         else:
             stmt = self.parse_expression_or_assignment()
+        if ts == 'async' and self.current_str() == 'def':
+            self.parse_error_at(self.current(),
+                                reason='Use --fast-parser to parse code using "async def"')
+            raise ParseError()
         if stmt is not None:
             stmt.set_line(t)
         return stmt, is_simple

mypy/parse.py

@@ -62,16 +62,16 @@
     ComparisonExpr, StarExpr, ARG_POS, ARG_NAMED, MroError, type_aliases,
     YieldFromExpr, NamedTupleExpr, NonlocalDecl,
     SetComprehension, DictionaryComprehension, TYPE_ALIAS, TypeAliasExpr,
-    YieldExpr, ExecStmt, Argument, BackquoteExpr, ImportBase, COVARIANT, CONTRAVARIANT,
+    YieldExpr, ExecStmt, Argument, BackquoteExpr, ImportBase, AwaitExpr,
     IntExpr, FloatExpr, UnicodeExpr,
-    INVARIANT, UNBOUND_IMPORTED
+    COVARIANT, CONTRAVARIANT, INVARIANT, UNBOUND_IMPORTED,
 )
 from mypy.visitor import NodeVisitor
 from mypy.traverser import TraverserVisitor
 from mypy.errors import Errors, report_internal_error
 from mypy.types import (
     NoneTyp, CallableType, Overloaded, Instance, Type, TypeVarType, AnyType,
-    FunctionLike, UnboundType, TypeList, ErrorType, TypeVarDef,
+    FunctionLike, UnboundType, TypeList, ErrorType, TypeVarDef, Void,
     replace_leading_arg_type, TupleType, UnionType, StarType, EllipsisType
 )
 from mypy.nodes import function_type, implicit_module_attrs
@@ -314,6 +314,13 @@ def visit_func_def(self, defn: FuncDef) -> None:
             # Second phase of analysis for function.
             self.errors.push_function(defn.name())
             self.analyze_function(defn)
+            if defn.is_coroutine and isinstance(defn.type, CallableType):
+                # A coroutine defined as `async def foo(...) -> T: ...`
+                # has external return type `Awaitable[T]`.
+                defn.type = defn.type.copy_modified(
+                    ret_type=Instance(
+                        self.named_type_or_none('typing.Awaitable').type,
+                        [defn.type.ret_type]))
             self.errors.pop_function()
 
     def prepare_method_signature(self, func: FuncDef) -> None:
@@ -1821,7 +1828,10 @@ def visit_yield_from_expr(self, e: YieldFromExpr) -> None:
         if not self.is_func_scope():  # not sure
             self.fail("'yield from' outside function", e, True, blocker=True)
         else:
-            self.function_stack[-1].is_generator = True
+            if self.function_stack[-1].is_coroutine:
+                self.fail("'yield from' in async function", e, True, blocker=True)
+            else:
+                self.function_stack[-1].is_generator = True
         if e.expr:
             e.expr.accept(self)
 
@@ -2074,10 +2084,20 @@ def visit_yield_expr(self, expr: YieldExpr) -> None:
         if not self.is_func_scope():
             self.fail("'yield' outside function", expr, True, blocker=True)
         else:
-            self.function_stack[-1].is_generator = True
+            if self.function_stack[-1].is_coroutine:
+                self.fail("'yield' in async function", expr, True, blocker=True)
+            else:
+                self.function_stack[-1].is_generator = True
         if expr.expr:
             expr.expr.accept(self)
 
+    def visit_await_expr(self, expr: AwaitExpr) -> None:
+        if not self.is_func_scope():
+            self.fail("'await' outside function", expr)
+        elif not self.function_stack[-1].is_coroutine:
+            self.fail("'await' outside coroutine ('async def')", expr)
+        expr.expr.accept(self)
+
     #
     # Helpers
     #

mypy/semanal.py

@@ -199,8 +199,10 @@ def visit_while_stmt(self, o):
         return self.dump(a, o)
 
     def visit_for_stmt(self, o):
-        a = [o.index]
-        a.extend([o.expr, o.body])
+        a = []
+        if o.is_async:
+            a.append(('Async', ''))
+        a.extend([o.index, o.expr, o.body])
         if o.else_body:
             a.append(('Else', o.else_body.body))
         return self.dump(a, o)
@@ -243,6 +245,9 @@ def visit_yield_from_stmt(self, o):
     def visit_yield_expr(self, o):
         return self.dump([o.expr], o)
 
+    def visit_await_expr(self, o):
+        return self.dump([o.expr], o)
+
     def visit_del_stmt(self, o):
         return self.dump([o.expr], o)
 
@@ -264,6 +269,8 @@ def visit_try_stmt(self, o):
 
     def visit_with_stmt(self, o):
         a = []
+        if o.is_async:
+            a.append(('Async', ''))
         for i in range(len(o.expr)):
             a.append(('Expr', [o.expr[i]]))
             if o.target[i]:

mypy/strconv.py

@@ -63,6 +63,7 @@
     'check-optional.test',
     'check-fastparse.test',
     'check-warnings.test',
+    'check-async-await.test',
 ]
 
 

mypy/test/testcheck.py

@@ -19,7 +19,7 @@
     ComparisonExpr, TempNode, StarExpr,
     YieldFromExpr, NamedTupleExpr, NonlocalDecl, SetComprehension,
     DictionaryComprehension, ComplexExpr, TypeAliasExpr, EllipsisExpr,
-    YieldExpr, ExecStmt, Argument, BackquoteExpr
+    YieldExpr, ExecStmt, Argument, BackquoteExpr, AwaitExpr,
 )
 from mypy.types import Type, FunctionLike, Instance
 from mypy.visitor import NodeVisitor
@@ -339,6 +339,9 @@ def visit_yield_from_expr(self, node: YieldFromExpr) -> Node:
     def visit_yield_expr(self, node: YieldExpr) -> Node:
         return YieldExpr(self.node(node.expr))
 
+    def visit_await_expr(self, node: AwaitExpr) -> Node:
+        return AwaitExpr(self.node(node.expr))
+
     def visit_call_expr(self, node: CallExpr) -> Node:
         return CallExpr(self.node(node.callee),
                         self.nodes(node.args),

mypy/treetransform.py

@@ -228,5 +228,8 @@ def visit_namedtuple_expr(self, o: 'mypy.nodes.NamedTupleExpr') -> T:
     def visit__promote_expr(self, o: 'mypy.nodes.PromoteExpr') -> T:
         pass
 
+    def visit_await_expr(self, o: 'mypy.nodes.AwaitExpr') -> T:
+        pass
+
     def visit_temp_node(self, o: 'mypy.nodes.TempNode') -> T:
         pass