Vast strides in accuracy for visit_await_expr().

Author: Guido van Rossum

mypy/checker.py

@@ -2019,30 +2019,24 @@ def visit_yield_expr(self, e: YieldExpr) -> Type:
         return self.get_generator_receive_type(return_type)
 
     def visit_await_expr(self, e: AwaitExpr) -> Type:
-        any_type = AnyType()
-        awaitable_type = self.named_generic_type('typing.Awaitable', [any_type])
-        generator_type = self.named_generic_type('typing.Generator', [any_type] * 3)
-        expected_type = UnionType([awaitable_type, generator_type])
+        expected_type = self.type_context[-1]
+        if expected_type is not None:
+            expected_type = self.named_generic_type('typing.Awaitable', [expected_type])
         actual_type = self.accept(e.expr, expected_type)
         if isinstance(actual_type, AnyType):
-            return any_type
-        if is_subtype(actual_type, generator_type):
-            if (isinstance(actual_type, Instance) and len(actual_type.args) == 3
-                    and actual_type.type.fullname() == 'typing.Generator'):
-                return actual_type.args[2]
-            else:
-                return any_type  # Must've been unparameterized Generator.
-        elif is_subtype(actual_type, awaitable_type):
+            return AnyType()
+        awaitable_type = self.named_generic_type('typing.Awaitable', [AnyType()])
+        if is_subtype(actual_type, awaitable_type):
             if (isinstance(actual_type, Instance) and len(actual_type.args) == 1
                     and actual_type.type.fullname() == 'typing.Awaitable'):
                 return actual_type.args[0]
             else:
-                return any_type  # Must've been unparameterized Awaitable.
-        msg = "{} (actual type {}, expected Awaitable or Generator)".format(
+                return AnyType()  # Must've been unparameterized Awaitable.
+        msg = "{} (actual type {}, expected Awaitable)".format(
             messages.INCOMPATIBLE_TYPES_IN_AWAIT,
             self.msg.format(actual_type))
         self.fail(msg, e)
-        return any_type
+        return AnyType()
 
     #
     # Helpers

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

@@ -21,16 +21,53 @@ async def f() -> int:
     return x
 [builtins fixtures/async_await.py]
 
-[case testAwaitGenerator]
+[case testAwaitDefaultContext]
 # options: fast_parser
-from typing import Any, Generator
-def g() -> Generator[int, int, int]:
-    x = yield 0
-    return x
+from typing import Any, TypeVar
+T = TypeVar('T')
+async def f(x: T) -> T:
+    y = await f(x)
+    reveal_type(y)
+    return y
+[out]
+main: note: In function "f":
+main:6: error: Revealed type is 'T`-1'
+
+[case testAwaitAnyContext]
+# options: fast_parser
+from typing import Any, TypeVar
+T = TypeVar('T')
+async def f(x: T) -> T:
+    y = await f(x)  # type: Any
+    reveal_type(y)
+    return y
+[out]
+main: note: In function "f":
+main:6: error: Revealed type is 'Any'
+
+[case testAwaitExplicitContext]
+# options: fast_parser
+from typing import Any, TypeVar
+T = TypeVar('T')
+async def f(x: T) -> T:
+    y = await f(x)  # type: int
+    reveal_type(y)
+[out]
+main: note: In function "f":
+main:5: error: Argument 1 to "f" has incompatible type "T"; expected "int"
+main:6: error: Revealed type is 'builtins.int'
+
+[case testAwaitGeneratorError]
+# options: fast_parser
+from typing import Any, Iterator
+def g() -> Iterator[Any]:
+    yield
 async def f() -> int:
     x = await g()
     return x
-[builtins fixtures/async_await.py]
+[out]
+main: note: In function "f":
+main: error: Incompatible types in await (actual type Iterator[Any], expected Awaitable)
 
 [case testAwaitArgumentError]
 # options: fast_parser
@@ -42,7 +79,7 @@ async def f() -> int:
 [builtins fixtures/async_await.py]
 [out]
 main: note: In function "f":
-main: error: Incompatible types in await (actual type "int", expected Awaitable or Generator)
+main: error: Incompatible types in await (actual type "int", expected Awaitable)
 
 [case testAwaitResultError]
 # options: fast_parser