From 2921c0ef9ee6d477fe5e098874f3188d944df208 Mon Sep 17 00:00:00 2001
From: Mikhail Golubev <mikhail.golubev@jetbrains.com>
Date: Mon, 12 Dec 2016 23:04:17 +0300
Subject: [PATCH] PY-20657 Loose return types of generators annotated to return
 not typing.Generator

Sometimes people use types other than typing.Generator to annotate result
types of generator functions, e.g. Iterable[T]. Because these other generic
types don't allow to specify type of expressions in return statements,
consider it to be Any, so as not to warn about mismatched actual and expected
return types in type checker inspection.
---
 .../com/jetbrains/python/psi/PyFunction.java  |  5 ++++
 .../inspections/PyTypeCheckerInspection.java  | 14 +++++++----
 .../python/psi/impl/PyFunctionImpl.java       | 24 +++++++++++++++++++
 .../GeneratorAnnotatedToReturnIterable.py     |  6 +++++
 .../Py3TypeCheckerInspectionTest.java         |  5 ++++
 5 files changed, 50 insertions(+), 4 deletions(-)
 create mode 100644 python/testData/inspections/PyTypeCheckerInspection/GeneratorAnnotatedToReturnIterable.py

diff --git a/python/psi-api/src/com/jetbrains/python/psi/PyFunction.java b/python/psi-api/src/com/jetbrains/python/psi/PyFunction.java
index 23399c5e79e2..fa668a445105 100644
--- a/python/psi-api/src/com/jetbrains/python/psi/PyFunction.java
+++ b/python/psi-api/src/com/jetbrains/python/psi/PyFunction.java
@@ -72,6 +72,11 @@ public interface PyFunction extends PsiNamedElement, StubBasedPsiElement<PyFunct
   @Nullable
   Modifier getModifier();
 
+  /**
+   * Checks whether the function contains a yield expression in its body.
+   */
+  boolean isGenerator();
+
   boolean isAsync();
 
   boolean isAsyncAllowed();
diff --git a/python/src/com/jetbrains/python/inspections/PyTypeCheckerInspection.java b/python/src/com/jetbrains/python/inspections/PyTypeCheckerInspection.java
index 0e43e4969937..be2502b27cf5 100644
--- a/python/src/com/jetbrains/python/inspections/PyTypeCheckerInspection.java
+++ b/python/src/com/jetbrains/python/inspections/PyTypeCheckerInspection.java
@@ -115,10 +115,16 @@ public class PyTypeCheckerInspection extends PyInspection {
         if (PyNames.FAKE_COROUTINE.equals(genericType.getName())) {
           return genericType.getIteratedItemType();
         }
-        else if (PyNames.FAKE_GENERATOR.equals(genericType.getName()) ||
-                 genericType instanceof PyClassType && "typing.Generator".equals(((PyClassType)genericType).getClassQName())) {
-          // Generator's type is parametrized as [YieldType, SendType, ReturnType]
-          return ContainerUtil.getOrElse(genericType.getElementTypes(myTypeEvalContext), 2, null);
+        else if (function.isGenerator()) {
+          if (PyNames.FAKE_GENERATOR.equals(genericType.getName()) ||
+              genericType instanceof PyClassType && "typing.Generator".equals(((PyClassType)genericType).getClassQName())) {
+            // Generator's type is parametrized as [YieldType, SendType, ReturnType]
+            return ContainerUtil.getOrElse(genericType.getElementTypes(myTypeEvalContext), 2, null);
+          }
+          else {
+            // Assume that any other return type annotation for a generator cannot contain its return type
+            return null;
+          }
         }
       }
 
diff --git a/python/src/com/jetbrains/python/psi/impl/PyFunctionImpl.java b/python/src/com/jetbrains/python/psi/impl/PyFunctionImpl.java
index 44e45ed23a39..682d22b3710f 100644
--- a/python/src/com/jetbrains/python/psi/impl/PyFunctionImpl.java
+++ b/python/src/com/jetbrains/python/psi/impl/PyFunctionImpl.java
@@ -693,6 +693,30 @@ public class PyFunctionImpl extends PyBaseElementImpl<PyFunctionStub> implements
     return null;
   }
 
+  @Override
+  public boolean isGenerator() {
+    final Ref<Boolean> result = new Ref<>(false);
+    getStatementList().accept(new PyRecursiveElementVisitor() {
+      @Override
+      public void visitPyYieldExpression(PyYieldExpression node) {
+        result.set(true);
+      }
+
+      @Override
+      public void visitPyFunction(PyFunction node) {
+        // Ignore nested functions
+      }
+
+      @Override
+      public void visitElement(PsiElement element) {
+        if (!result.get()) {
+          super.visitElement(element);
+        }
+      }
+    });
+    return result.get();
+  }
+
   @Override
   public boolean isAsync() {
     final PyFunctionStub stub = getStub();
diff --git a/python/testData/inspections/PyTypeCheckerInspection/GeneratorAnnotatedToReturnIterable.py b/python/testData/inspections/PyTypeCheckerInspection/GeneratorAnnotatedToReturnIterable.py
new file mode 100644
index 000000000000..0d6b9f8fba81
--- /dev/null
+++ b/python/testData/inspections/PyTypeCheckerInspection/GeneratorAnnotatedToReturnIterable.py
@@ -0,0 +1,6 @@
+from typing import Iterable
+
+
+def f() -> Iterable[int]:
+    for i in range(10):
+        yield i
\ No newline at end of file
diff --git a/python/testSrc/com/jetbrains/python/inspections/Py3TypeCheckerInspectionTest.java b/python/testSrc/com/jetbrains/python/inspections/Py3TypeCheckerInspectionTest.java
index a4f060ca6903..554a3387219e 100644
--- a/python/testSrc/com/jetbrains/python/inspections/Py3TypeCheckerInspectionTest.java
+++ b/python/testSrc/com/jetbrains/python/inspections/Py3TypeCheckerInspectionTest.java
@@ -179,6 +179,11 @@ public class Py3TypeCheckerInspectionTest extends PyTestCase {
     doTest();
   }
 
+  // PY-20657
+  public void testGeneratorAnnotatedToReturnIterable() {
+    doTest();
+  }
+
   // PY-21083
   public void testFloatFromhex() {
     doTest();