PY-18803 Annotate Types inserts return type annotation over the old one, not after

This commit is contained in:
Mikhail Golubev
2016-03-17 19:11:11 +03:00
committed by Andrey Vlasovskikh
parent 243c730c82
commit acbf0524e5
@@ -137,8 +137,7 @@ public class SpecifyTypeInPy3AnnotationsIntention extends TypeIntention {
static String returnType(@NotNull PyFunction function) {
String returnType = PyNames.OBJECT;
final PySignature signature = PySignatureCacheManager.getInstance(function.getProject()).findSignature(
function);
final PySignature signature = PySignatureCacheManager.getInstance(function.getProject()).findSignature(function);
if (signature != null) {
returnType = ObjectUtils.chooseNotNull(signature.getReturnTypeQualifiedName(), returnType);
}
@@ -148,22 +147,28 @@ public class SpecifyTypeInPy3AnnotationsIntention extends TypeIntention {
public static PyExpression annotateReturnType(Project project, PyFunction function, boolean createTemplate) {
String returnType = returnType(function);
final String annotationText = " -> " + returnType;
final PsiElement prevElem = PyPsiUtils.getPrevNonCommentSibling(function.getStatementList(), true);
assert prevElem != null;
final String annotationText = "-> " + returnType;
final PsiDocumentManager manager = PsiDocumentManager.getInstance(project);
Document documentWithCallable = manager.getDocument(function.getContainingFile());
if (documentWithCallable != null) {
try {
final TextRange range = prevElem.getTextRange();
manager.doPostponedOperationsAndUnblockDocument(documentWithCallable);
if (prevElem.getNode().getElementType() == PyTokenTypes.COLON) {
documentWithCallable.insertString(range.getStartOffset(), annotationText);
final PyAnnotation oldAnnotation = function.getAnnotation();
if (oldAnnotation != null) {
final TextRange oldRange = oldAnnotation.getTextRange();
documentWithCallable.replaceString(oldRange.getStartOffset(), oldRange.getEndOffset(), annotationText);
}
else {
documentWithCallable.insertString(range.getEndOffset(), annotationText + ":");
final PsiElement prevElem = PyPsiUtils.getPrevNonCommentSibling(function.getStatementList(), true);
assert prevElem != null;
final TextRange range = prevElem.getTextRange();
if (prevElem.getNode().getElementType() == PyTokenTypes.COLON) {
documentWithCallable.insertString(range.getStartOffset(), " " + annotationText);
}
else {
documentWithCallable.insertString(range.getEndOffset(), " " + annotationText + ":");
}
}
}
finally {