Fix crash when checking async generator return type (#5123)

Author: ilevkivskyi

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)
 

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]