public static PyExpression annotateReturnType(Project project, PyFunction function, boolean createTemplate) {
    String returnType = returnType(function);

    final String annotationText = "-> " + returnType;

    PyFunction annotatedFunction = PyUtil.updateDocumentUnblockedAndCommitted(function, document -> {
      final PyAnnotation oldAnnotation = function.getAnnotation();
      if (oldAnnotation != null) {
        final TextRange oldRange = oldAnnotation.getTextRange();
        document.replaceString(oldRange.getStartOffset(), oldRange.getEndOffset(), annotationText);
      }
      else {
        final PsiElement prevElem = PyPsiUtils.getPrevNonCommentSibling(function.getStatementList(), true);
        assert prevElem != null;
        final TextRange range = prevElem.getTextRange();
        if (prevElem.getNode().getElementType() == PyTokenTypes.COLON) {
          document.insertString(range.getStartOffset(), " " + annotationText);
        }
        else {
          document.insertString(range.getEndOffset(), " " + annotationText + ":");
        }
      }
      return CodeInsightUtilCore.forcePsiPostprocessAndRestoreElement(function);
    });

    if (annotatedFunction == null) {
      return null;
    }

    final PyAnnotation annotation = annotatedFunction.getAnnotation();
    assert annotation != null;
    final PyExpression annotationValue = annotation.getValue();
    assert annotationValue != null : "Generated function must have annotation";

    if (createTemplate) {
      final int offset = annotationValue.getTextOffset();

      final TemplateBuilder builder = TemplateBuilderFactory.getInstance().createTemplateBuilder(annotationValue);
      builder.replaceRange(TextRange.create(0, returnType.length()), returnType);
      final Editor targetEditor = PythonUiService.getInstance().openTextEditor(project, annotatedFunction.getContainingFile().getVirtualFile(), offset);
      if (targetEditor != null) {
        builder.run(targetEditor, true);
      }
    }
    return annotationValue;
  }
public static PyExpression annotateReturnType(Project project, PyFunction function, boolean createTemplate)