From 75d6657d79ce6c1f30c957d20bd7884a179e7fbd Mon Sep 17 00:00:00 2001
From: Ivan Levkivskyi <ilevkivskyi@dropbox.com>
Date: Wed, 30 May 2018 22:44:12 +0100
Subject: [PATCH] Fix crash when checking async generator return type

---
 mypy/checker.py                       | 10 ++++++---
 test-data/unit/check-async-await.test | 30 +++++++++++++++++++++++++++
 2 files changed, 37 insertions(+), 3 deletions(-)

diff --git a/mypy/checker.py b/mypy/checker.py
index 9e19bf9d96da..7875672bd59b 100644
--- a/mypy/checker.py
+++ b/mypy/checker.py
@@ -929,9 +929,13 @@ def is_unannotated_any(t: Type) -> bool:
                 ret_type = fdef.type.ret_type
                 if is_unannotated_any(ret_type):
                     self.fail(messages.RETURN_TYPE_EXPECTED, fdef)
-                elif (fdef.is_coroutine and isinstance(ret_type, Instance) and
-                      is_unannotated_any(self.get_coroutine_return_type(ret_type))):
-                    self.fail(messages.RETURN_TYPE_EXPECTED, fdef)
+                elif fdef.is_generator:
+                    if is_unannotated_any(self.get_generator_return_type(ret_type,
+                                                                        fdef.is_coroutine)):
+                        self.fail(messages.RETURN_TYPE_EXPECTED, fdef)
+                elif fdef.is_coroutine and isinstance(ret_type, Instance):
+                    if is_unannotated_any(self.get_coroutine_return_type(ret_type)):
+                        self.fail(messages.RETURN_TYPE_EXPECTED, fdef)
                 if any(is_unannotated_any(t) for t in fdef.type.arg_types):
                     self.fail(messages.ARGUMENT_TYPE_EXPECTED, fdef)
 
diff --git a/test-data/unit/check-async-await.test b/test-data/unit/check-async-await.test
index 62222ce02e8b..09e294031f89 100644
--- a/test-data/unit/check-async-await.test
+++ b/test-data/unit/check-async-await.test
@@ -669,3 +669,33 @@ async def decorated_host_coroutine() -> None:
 [builtins fixtures/async_await.pyi]
 [typing fixtures/typing-full.pyi]
 [out]
+
+[case testAsyncGenDisallowUntyped]
+# flags: --disallow-untyped-defs
+# These should not crash
+from typing import AsyncGenerator, Any
+
+async def f() -> AsyncGenerator[int, None]:
+    yield 0
+
+async def g() -> AsyncGenerator[Any, None]:
+    yield 0
+[builtins fixtures/async_await.pyi]
+[typing fixtures/typing-full.pyi]
+[out]
+
+[case testAsyncGenDisallowUntypedTriggers]
+# flags: --disallow-untyped-defs
+from typing import AsyncGenerator, Any
+
+async def f() -> AsyncGenerator[Any, Any]:
+    yield None
+
+async def h() -> Any:
+    yield 0
+
+async def g():  # E: Function is missing a type annotation
+    yield 0
+[builtins fixtures/async_await.pyi]
+[typing fixtures/typing-full.pyi]
+[out]