From acbf0524e5f369dcaf53124ccbb19fc43c6b0fec Mon Sep 17 00:00:00 2001
From: Mikhail Golubev <mikhail.golubev@jetbrains.com>
Date: Thu, 17 Mar 2016 17:43:52 +0300
Subject: [PATCH] PY-18803 Annotate Types inserts return type annotation over
 the old one, not after

---
 .../SpecifyTypeInPy3AnnotationsIntention.java | 25 +++++++++++--------
 1 file changed, 15 insertions(+), 10 deletions(-)

diff --git a/python/src/com/jetbrains/python/codeInsight/intentions/SpecifyTypeInPy3AnnotationsIntention.java b/python/src/com/jetbrains/python/codeInsight/intentions/SpecifyTypeInPy3AnnotationsIntention.java
index f9b41ab1c7b7..a2b2735c5aa1 100644
--- a/python/src/com/jetbrains/python/codeInsight/intentions/SpecifyTypeInPy3AnnotationsIntention.java
+++ b/python/src/com/jetbrains/python/codeInsight/intentions/SpecifyTypeInPy3AnnotationsIntention.java
@@ -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 {