From 46cff81ceedb75075280c8415b45fc67f376d058 Mon Sep 17 00:00:00 2001
From: Mikhail Golubev <mikhail.golubev@jetbrains.com>
Date: Thu, 18 Jan 2018 14:39:25 +0300
Subject: [PATCH] PY-16412 SpecifyTypeInPy3AnnotationIntention uses
 PyUtil#updateDocument

---
 .../SpecifyTypeInPy3AnnotationsIntention.java | 51 ++++++++-----------
 1 file changed, 20 insertions(+), 31 deletions(-)

diff --git a/python/src/com/jetbrains/python/codeInsight/intentions/SpecifyTypeInPy3AnnotationsIntention.java b/python/src/com/jetbrains/python/codeInsight/intentions/SpecifyTypeInPy3AnnotationsIntention.java
index 18ff6b2d9640..5ee79e4c735e 100644
--- a/python/src/com/jetbrains/python/codeInsight/intentions/SpecifyTypeInPy3AnnotationsIntention.java
+++ b/python/src/com/jetbrains/python/codeInsight/intentions/SpecifyTypeInPy3AnnotationsIntention.java
@@ -17,14 +17,12 @@ package com.jetbrains.python.codeInsight.intentions;
 
 import com.intellij.codeInsight.CodeInsightUtilCore;
 import com.intellij.codeInsight.template.*;
-import com.intellij.openapi.editor.Document;
 import com.intellij.openapi.editor.Editor;
 import com.intellij.openapi.fileEditor.FileEditorManager;
 import com.intellij.openapi.fileEditor.OpenFileDescriptor;
 import com.intellij.openapi.project.Project;
 import com.intellij.openapi.util.TextRange;
 import com.intellij.openapi.util.text.StringUtil;
-import com.intellij.psi.PsiDocumentManager;
 import com.intellij.psi.PsiElement;
 import com.intellij.psi.PsiFile;
 import com.intellij.psi.PsiReference;
@@ -153,36 +151,31 @@ public class SpecifyTypeInPy3AnnotationsIntention extends TypeIntention {
 
     final String annotationText = "-> " + returnType;
 
-    final PsiDocumentManager manager = PsiDocumentManager.getInstance(project);
-    Document documentWithCallable = manager.getDocument(function.getContainingFile());
-    if (documentWithCallable != null) {
-      try {
-        manager.doPostponedOperationsAndUnblockDocument(documentWithCallable);
-        final PyAnnotation oldAnnotation = function.getAnnotation();
-        if (oldAnnotation != null) {
-          final TextRange oldRange = oldAnnotation.getTextRange();
-          documentWithCallable.replaceString(oldRange.getStartOffset(), oldRange.getEndOffset(), annotationText);
+    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 {
-          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 + ":");
-          }
+          document.insertString(range.getEndOffset(), " " + annotationText + ":");
         }
       }
-      finally {
-        manager.commitDocument(documentWithCallable);
-      }
+      return CodeInsightUtilCore.forcePsiPostprocessAndRestoreElement(function);
+    });
+
+    if (annotatedFunction == null) {
+      return null;
     }
 
-
-    function = CodeInsightUtilCore.forcePsiPostprocessAndRestoreElement(function);
-    final PyAnnotation annotation = function.getAnnotation();
+    final PyAnnotation annotation = annotatedFunction.getAnnotation();
     assert annotation != null;
     final PyExpression annotationValue = annotation.getValue();
     assert annotationValue != null : "Generated function must have annotation";
@@ -193,11 +186,7 @@ public class SpecifyTypeInPy3AnnotationsIntention extends TypeIntention {
       final TemplateBuilder builder = TemplateBuilderFactory.getInstance().createTemplateBuilder(annotationValue);
       builder.replaceRange(TextRange.create(0, returnType.length()), returnType);
       final Template template = ((TemplateBuilderImpl)builder).buildInlineTemplate();
-      final OpenFileDescriptor descriptor = new OpenFileDescriptor(
-        project,
-        function.getContainingFile().getVirtualFile(),
-        offset
-      );
+      final OpenFileDescriptor descriptor = new OpenFileDescriptor(project, annotatedFunction.getContainingFile().getVirtualFile(), offset);
       final Editor targetEditor = FileEditorManager.getInstance(project).openTextEditor(descriptor, true);
       if (targetEditor != null) {
         targetEditor.getCaretModel().moveToOffset(offset);