Generalize yield statement function type

mypy/checker.py

@@ -1458,10 +1458,18 @@ def count_nested_types(self, typ: Instance, check_type: str) -> int:
     def visit_yield_stmt(self, s: YieldStmt) -> Type:
         return_type = self.return_types[-1]
         if isinstance(return_type, Instance):
-            if return_type.type.fullname() != 'typing.Iterator':
+            if not is_subtype(self.named_generic_type('typing.Generator',
+                                                      [AnyType(), AnyType(), AnyType()]),
+                              return_type):
                 self.fail(messages.INVALID_RETURN_TYPE_FOR_YIELD, s)
                 return None
-            expected_item_type = return_type.args[0]
+            if return_type.args:
+                expected_item_type = return_type.args[0]
+            else:
+                # if the declared supertype of `Generator` has no type
+                # parameters (i.e. is `object`), then the yielded values can't
+                # be accessed so any type is acceptable.
+                expected_item_type = AnyType()
         elif isinstance(return_type, AnyType):
             expected_item_type = AnyType()
         else:

mypy/messages.py

@@ -30,7 +30,7 @@
 INVALID_EXCEPTION = 'Exception must be derived from BaseException'
 INVALID_EXCEPTION_TYPE = 'Exception type must be derived from BaseException'
 INVALID_RETURN_TYPE_FOR_YIELD = \
-    'Iterator function return type expected for "yield"'
+    'The return type of a generator function should be "Generator" or one of its supertypes'
 INVALID_RETURN_TYPE_FOR_YIELD_FROM = \
     'Iterable function return type expected for "yield from"'
 INCOMPATIBLE_TYPES = 'Incompatible types'

mypy/test/data/check-statements.test

@@ -647,6 +647,26 @@ def f() -> Iterator[int]:
 [out]
 main: note: In function "f":
 
+[case testYieldInFunctionReturningGenerator]
+from typing import Generator
+def f() -> Generator[int, None, None]:
+    yield 1
+[builtins fixtures/for.py]
+[out]
+
+[case testYieldInFunctionReturningIterable]
+from typing import Iterable
+def f() -> Iterable[int]:
+    yield 1
+[builtins fixtures/for.py]
+[out]
+
+[case testYieldInFunctionReturningObject]
+def f() -> object:
+    yield 1
+[builtins fixtures/for.py]
+[out]
+
 [case testYieldInFunctionReturningAny]
 from typing import Any
 def f() -> Any:
@@ -656,7 +676,7 @@ def f() -> Any:
 [case testYieldInFunctionReturningFunction]
 from typing import Callable
 def f() -> Callable[[], None]:
-    yield object() # E: Iterator function return type expected for "yield"
+    yield object() # E: The return type of a generator function should be "Generator" or one of its supertypes
 [out]
 main: note: In function "f":
 
@@ -668,7 +688,7 @@ def f():
 [case testWithInvalidInstanceReturnType]
 import typing
 def f() -> int:
-    yield 1 # E: Iterator function return type expected for "yield"
+    yield 1 # E: The return type of a generator function should be "Generator" or one of its supertypes
 [builtins fixtures/for.py]
 [out]
 main: note: In function "f":

mypy/test/data/fixtures/for.py

@@ -1,6 +1,6 @@
 # builtins stub used in for statement test cases
 
-from typing import TypeVar, Generic, Iterable, Iterator
+from typing import TypeVar, Generic, Iterable, Iterator, Generator
 from abc import abstractmethod, ABCMeta
 
 t = TypeVar('t')

mypy/test/data/lib-stub/typing.py

@@ -23,6 +23,8 @@
 Set = 0
 
 T = TypeVar('T')
+U = TypeVar('U')
+V = TypeVar('V')
 
 class Container(Generic[T]):
     @abstractmethod
@@ -41,6 +43,16 @@ class Iterator(Iterable[T], Generic[T]):
     @abstractmethod
     def __next__(self) -> T: pass
 
+class Generator(Iterator[T], Generic[T, U, V]):
+    @abstractmethod
+    def send(self, value: U) -> T: pass
+
+    @abstractmethod
+    def throw(self, typ: Any, val: Any=None, tb=None) -> None: pass
+
+    @abstractmethod
+    def close(self) -> None: pass
+
 class Sequence(Generic[T]):
     @abstractmethod
     def __getitem__(self, n: Any) -> T: pass