Repository: quora/pyanalyze
Pull request #756 (2024-03-26 15:58:07)

Detect incorrect generator return types


Patch:
diff --git a/docs/changelog.md b/docs/changelog.md
index 87ee718..02ebc32 100644
--- a/docs/changelog.md
+++ b/docs/changelog.md
@@ -2,13 +2,16 @@
 
 ## Unreleased
 
+- New error code `generator_return` is raised when a generator does not
+  return an iterable type, or an async generator does not return an async
+  iterable type (#756)
 - Fix type narrowing for certain conditionals using `or` (#755)
 - Fix incorrect `undefined_name` errors when a class is nested in a nested
   function and uses a name from the outer function (#750)
 - Fix incorrect `possibly_undefined_name` error on certain uses of the
   walrus operator (#749)
 - Fix narrowing on `isinstance` calls with arguments that are not
-  instance of `typing`, such as unions and certain typing special forms (#747)
+  instances of `type`, such as unions and certain typing special forms (#747)
 - Detect invalid calls to `isinstance` (#747)
 - Support calls to `TypeVar` and several other typing constructs in
   code that is not executed (e.g., under `if TYPE_CHECKING`) (#746)
diff --git a/pyanalyze/error_code.py b/pyanalyze/error_code.py
index 4c781ea..83d92cb 100644
--- a/pyanalyze/error_code.py
+++ b/pyanalyze/error_code.py
@@ -167,6 +167,7 @@ ErrorCode = ErrorRegistry(
         ),
         Error("invalid_typeguard", "Invalid use of TypeGuard or TypeIs"),
         Error("readonly_typeddict", "TypedDict is read-only"),
+        Error("generator_return", "Generator must return an iterable"),
     ]
 )
 
diff --git a/pyanalyze/name_check_visitor.py b/pyanalyze/name_check_visitor.py
index 40370b6..6ca7e6a 100644
--- a/pyanalyze/name_check_visitor.py
+++ b/pyanalyze/name_check_visitor.py
@@ -2049,6 +2049,33 @@ class NameCheckVisitor(node_visitor.ReplacingNodeVisitor):
                         node, error_code=ErrorCode.override_does_not_override
                     )
 
+        if info.return_annotation is not None:
+            assert node.returns is not None
+            if result.is_generator:
+                if isinstance(node, ast.FunctionDef):
+                    if info.async_kind is AsyncFunctionKind.non_async:
+                        can_assign = TypedValue(collections.abc.Iterable).can_assign(
+                            info.return_annotation, self
+                        )
+                        if isinstance(can_assign, CanAssignError):
+                            self._show_error_if_checking(
+                                node.returns,
+                                "Generator function must return an iterable",
+                                error_code=ErrorCode.generator_return,
+                                detail=can_assign.display(),
+                            )
+                else:
+                    can_assign = TypedValue(collections.abc.AsyncIterable).can_assign(
+                        info.return_annotation, self
+                    )
+                    if isinstance(can_assign, CanAssignError):
+                        self._show_error_if_checking(
+                            node.returns,
+                            "Async generator function must return an async iterable",
+                            error_code=ErrorCode.generator_return,
+                            detail=can_assign.display(),
+                        )
+
         if node.returns is None:
             if (
                 result.has_return


Test patch:
diff --git a/pyanalyze/test_generators.py b/pyanalyze/test_generators.py
index 82ae2d5..2101f4b 100644
--- a/pyanalyze/test_generators.py
+++ b/pyanalyze/test_generators.py
@@ -43,3 +43,43 @@ class TestGenerator(TestNameCheckVisitorBase):
         def caller() -> Iterable[int]:
             x = yield from gen(True)
             assert_is_value(x, AnyValue(AnySource.generic_argument))
+
+
+class TestGeneratorReturn(TestNameCheckVisitorBase):
+    @assert_passes()
+    def test_sync(self):
+        from typing import Generator, Iterable
+
+        from asynq import ConstFuture, asynq
+
+        def gen() -> int:  # E: generator_return
+            yield 1
+
+        def caller() -> int:  # E: generator_return
+            x = yield from [1, 2]
+            print(x)
+
+        def gen2() -> Iterable[int]:
+            yield 1
+
+        def caller2() -> Generator[int, None, None]:
+            x = yield from [1, 2]
+            print(x)
+
+        @asynq()
+        def asynq_gen() -> int:
+            x = yield ConstFuture(1)
+            return x
+
+    @assert_passes()
+    def test_async(self):
+        from typing import AsyncGenerator, AsyncIterable
+
+        async def gen() -> int:  # E: generator_return
+            yield 1
+
+        async def gen2() -> AsyncIterable[int]:
+            yield 1
+
+        async def gen3() -> AsyncGenerator[int, None]:
+            yield 1