PyExpression (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; }