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 implements return null; } + @Override + public boolean isGenerator() { + final Ref 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();