@Override
  public boolean isGenerator() {
    final PyFunctionStub stub = getStub();
    if (stub != null) {
      return stub.isGenerator();
    }

    Boolean result = myIsGenerator;
    if (result == null) {
      Ref<Boolean> containsYield = Ref.create(false);
      getStatementList().accept(new PyRecursiveElementVisitor() {
        @Override
        public void visitPyYieldExpression(@NotNull PyYieldExpression node) {
          containsYield.set(true);
        }

        @Override
        public void visitPyFunction(@NotNull PyFunction node) {
          // Ignore nested functions
        }

        @Override
        public void visitElement(@NotNull PsiElement element) {
          if (!containsYield.get()) {
            super.visitElement(element);
          }
        }
      });
      myIsGenerator = result = containsYield.get();
    }
    return result;
  }
@Override
  public boolean isGenerator()