Rewrite Generator[..., None, None] to Iterator[...] (#110) (Fixes #4)

monkeytype/typing.py

@@ -242,8 +242,19 @@ def rewrite(self, typ):
         return typ
 
 
+class RewriteGenerator(TypeRewriter):
+    """Returns an Iterator, if the send_type and return_type of a Generator is None"""
+
+    def rewrite_Generator(self, typ):
+        args = typ.__args__
+        if args[1] is NoneType and args[2] is NoneType:
+            return Iterator[args[0]]
+        return typ
+
+
 DEFAULT_REWRITER = ChainedRewriter((
     RemoveEmptyContainers(),
     RewriteConfigDict(),
     RewriteLargeUnion(),
+    RewriteGenerator(),
 ))

monkeytype/typing.pyi

@@ -71,4 +71,8 @@ class NoOpRewriter(TypeRewriter):
     ...
 
 
+class RewriteGenerator(TypeRewriter):
+    ...
+
+
 DEFAULT_REWRITER: TypeRewriter = ...

tests/test_typing.py

@@ -14,6 +14,7 @@
     Tuple,
     Type,
     Union,
+    Generator,
 )
 
 import pytest
@@ -26,7 +27,9 @@
     get_type,
     get_type_str,
     shrink_types,
+    RewriteGenerator,
 )
+
 from .util import Dummy
 
 
@@ -209,3 +212,22 @@ class G(C):
     def test_rewrite(self, typ, expected):
         rewritten = RewriteLargeUnion(2).rewrite(typ)
         assert rewritten == expected
+
+
+class TestRewriteGenerator:
+    @pytest.mark.parametrize(
+        'typ, expected',
+        [
+            # Should not rewrite
+            (Generator[int, None, int], Generator[int, None, int]),
+
+            # Should not rewrite
+            (Generator[int, int, None], Generator[int, int, None]),
+
+            # Should rewrite to Iterator[int]
+            (Generator[int, NoneType, NoneType], Iterator[int])
+        ],
+    )
+    def test_rewrite(self, typ, expected):
+        rewritten = RewriteGenerator().rewrite(typ)
+        assert rewritten == expected