From 541f8ca9e67e817596ddf6a44fef9181246c2dee Mon Sep 17 00:00:00 2001
From: Mikhail Golubev <mikhail.golubev@jetbrains.com>
Date: Wed, 29 Oct 2014 13:51:36 +0300
Subject: [PATCH] PY-13297 Use only PSI means to insert new line after comment

---
 .../inspections/quickfix/AddEncodingQuickFix.java    | 12 ++++++------
 ...codingInEmptyFile.py => AddEncodingAtLastLine.py} |  2 +-
 ...yFile_after.py => AddEncodingAtLastLine_after.py} |  1 +
 .../testSrc/com/jetbrains/python/PyQuickFixTest.java |  4 ++--
 4 files changed, 10 insertions(+), 9 deletions(-)
 rename python/testData/inspections/{AddEncodingInEmptyFile.py => AddEncodingAtLastLine.py} (74%)
 rename python/testData/inspections/{AddEncodingInEmptyFile_after.py => AddEncodingAtLastLine_after.py} (50%)

diff --git a/python/src/com/jetbrains/python/inspections/quickfix/AddEncodingQuickFix.java b/python/src/com/jetbrains/python/inspections/quickfix/AddEncodingQuickFix.java
index 82f61cfb1b57..3e3e55028abe 100644
--- a/python/src/com/jetbrains/python/inspections/quickfix/AddEncodingQuickFix.java
+++ b/python/src/com/jetbrains/python/inspections/quickfix/AddEncodingQuickFix.java
@@ -25,9 +25,9 @@ import com.intellij.openapi.fileEditor.FileEditorManager;
 import com.intellij.openapi.fileEditor.TextEditor;
 import com.intellij.openapi.project.Project;
 import com.intellij.psi.PsiComment;
-import com.intellij.psi.PsiDocumentManager;
 import com.intellij.psi.PsiElement;
 import com.intellij.psi.PsiFile;
+import com.intellij.psi.PsiWhiteSpace;
 import com.jetbrains.python.PyBundle;
 import com.jetbrains.python.inspections.PyEncodingUtil;
 import com.jetbrains.python.psi.LanguageLevel;
@@ -71,20 +71,20 @@ public class AddEncodingQuickFix implements LocalQuickFix {
     if (firstLine instanceof PsiComment && firstLine.getText().startsWith("#!")) {
       firstLine = firstLine.getNextSibling();
     }
+    final LanguageLevel languageLevel = LanguageLevel.forElement(file);
     final String commentText = String.format(PyEncodingUtil.ENCODING_FORMAT_PATTERN[myEncodingFormatIndex], myDefaultEncoding);
     final PyElementGenerator elementGenerator = PyElementGenerator.getInstance(project);
-    PsiComment encodingComment = elementGenerator.createFromText(LanguageLevel.forElement(file), PsiComment.class, commentText);
+    PsiComment encodingComment = elementGenerator.createFromText(languageLevel, PsiComment.class, commentText);
     encodingComment = (PsiComment)file.addBefore(encodingComment, firstLine);
 
     final FileEditor fileEditor = FileEditorManager.getInstance(project).getSelectedEditor(element.getContainingFile().getVirtualFile());
     if (fileEditor instanceof TextEditor) {
+      if (encodingComment.getNextSibling() == null || !encodingComment.getNextSibling().textContains('\n')) {
+        file.addAfter(elementGenerator.createFromText(languageLevel, PsiWhiteSpace.class, "\n"), encodingComment);
+      }
       final Editor editor = ((TextEditor)fileEditor).getEditor();
       final Document document = editor.getDocument();
       final int insertedLineNumber = document.getLineNumber(encodingComment.getTextOffset());
-      if (insertedLineNumber == document.getLineCount() - 1) {
-        PsiDocumentManager.getInstance(project).doPostponedOperationsAndUnblockDocument(document);
-        document.insertString(document.getLineEndOffset(insertedLineNumber), "\n");
-      }
       editor.getCaretModel().moveToLogicalPosition(new LogicalPosition(insertedLineNumber + 1, 0));
     }
   }
diff --git a/python/testData/inspections/AddEncodingInEmptyFile.py b/python/testData/inspections/AddEncodingAtLastLine.py
similarity index 74%
rename from python/testData/inspections/AddEncodingInEmptyFile.py
rename to python/testData/inspections/AddEncodingAtLastLine.py
index 376db9193fc4..a7868cf286f9 100644
--- a/python/testData/inspections/AddEncodingInEmptyFile.py
+++ b/python/testData/inspections/AddEncodingAtLastLine.py
@@ -1 +1 @@
-<warning descr="No encoding specified for file"><caret></warning>
+<warning descr="No encoding specified for file"><caret></warning>#!/usr/bin/env python
\ No newline at end of file
diff --git a/python/testData/inspections/AddEncodingInEmptyFile_after.py b/python/testData/inspections/AddEncodingAtLastLine_after.py
similarity index 50%
rename from python/testData/inspections/AddEncodingInEmptyFile_after.py
rename to python/testData/inspections/AddEncodingAtLastLine_after.py
index 76ea4ea7c759..5e9beaa6da1d 100644
--- a/python/testData/inspections/AddEncodingInEmptyFile_after.py
+++ b/python/testData/inspections/AddEncodingAtLastLine_after.py
@@ -1,2 +1,3 @@
+#!/usr/bin/env python
 # coding=utf-8
 <caret>
\ No newline at end of file
diff --git a/python/testSrc/com/jetbrains/python/PyQuickFixTest.java b/python/testSrc/com/jetbrains/python/PyQuickFixTest.java
index c7e8d138aa83..6cc2e578fb30 100644
--- a/python/testSrc/com/jetbrains/python/PyQuickFixTest.java
+++ b/python/testSrc/com/jetbrains/python/PyQuickFixTest.java
@@ -356,8 +356,8 @@ public class PyQuickFixTest extends PyTestCase {
   }
 
   // PY-13297
-  public void testAddEncodingInEmptyFile() {
-    doInspectionTest("AddEncodingInEmptyFile.py", PyMandatoryEncodingInspection.class,
+  public void testAddEncodingAtLastLine() {
+    doInspectionTest("AddEncodingAtLastLine.py", PyMandatoryEncodingInspection.class,
                      PyBundle.message("QFIX.add.encoding"), true, true);
   }