PY-17564 Use AddImportHelper#addFromImport in MoveFromFutureImportQuickfix

It's safer when PyBlock.IMPORT_GROUP_BEGIN marker is inserted for new
imports in single place (via methods of AddImportHelper).

As an additional benefit of this approach, moved __from__ import
statement is inserted in alphabetical order.

diff --git a/python/src/com/jetbrains/python/codeInsight/imports/AddImportHelper.java b/python/src/com/jetbrains/python/codeInsight/imports/AddImportHelper.java
index c00cf762e446a5f016ea27aeed0fdad40184770e..cf77aa8668d7d8e4568497c6fca91ca7bdacc1ec 100644 (file)
--- a/python/src/com/jetbrains/python/codeInsight/imports/AddImportHelper.java
+++ b/python/src/com/jetbrains/python/codeInsight/imports/AddImportHelper.java
@@ -311,7 +311,7 @@ public class AddImportHelper {
   }
 
   /**
-   * Adds a new {@link PyFromImportStatement} statement below other top-level imports or as specified by anchor.
+   * Adds a new {@link PyFromImportStatement} statement within other top-level imports or as specified by anchor.
    *
    * @param file   where to operate
    * @param from   import source (reference after {@code from} keyword)
@@ -330,18 +330,37 @@ public class AddImportHelper {
                                             @Nullable PsiElement anchor) {
     final PyElementGenerator generator = PyElementGenerator.getInstance(file.getProject());
     final LanguageLevel languageLevel = LanguageLevel.forElement(file);
-    final PyFromImportStatement nodeToInsert = generator.createFromImportStatement(languageLevel, from, name, asName);
+    final PyFromImportStatement newImport = generator.createFromImportStatement(languageLevel, from, name, asName);
+    addFromImportStatement(file, newImport, priority, anchor);
+  }
+
+  /**
+   * Adds a new {@link PyFromImportStatement} statement within other top-level imports or as specified by anchor.
+   *
+   * @param file      where to operate
+   * @param newImport new "from import" statement to insert. It may be generated, because it won't be used for resolving anyway. 
+   *                  You might want to use overloaded version of this method to generate such statement automatically.
+   * @param anchor    place where the imported name was used. It will be used to determine proper block where new import should be inserted,
+   *                  e.g. inside conditional block or try/except statement. Also if anchor is another import statement, new import statement
+   *                  will be inserted right after it.
+   * @see #addFromImportStatement(PsiFile, String, String, String, ImportPriority, PsiElement)
+   * @see #addFromImportStatement
+   */
+  public static void addFromImportStatement(@NotNull PsiFile file,
+                                            @NotNull PyFromImportStatement newImport,
+                                            @Nullable ImportPriority priority,
+                                            @Nullable PsiElement anchor) {
     try {
-      final PyImportStatementBase importStatement = PsiTreeUtil.getParentOfType(anchor, PyImportStatementBase.class, false);
+      final PyImportStatementBase parentImport = PsiTreeUtil.getParentOfType(anchor, PyImportStatementBase.class, false);
       final PsiElement insertParent;
-      if (importStatement != null && importStatement.getContainingFile() == file) {
-        insertParent = importStatement.getParent();
+      if (parentImport != null && parentImport.getContainingFile() == file) {
+        insertParent = parentImport.getParent();
       }
       else {
         insertParent = file;
       }
       if (InjectedLanguageManager.getInstance(file.getProject()).isInjectedFragment(file)) {
-        final PsiElement element = insertParent.addBefore(nodeToInsert, getInsertPosition(insertParent, nodeToInsert, priority));
+        final PsiElement element = insertParent.addBefore(newImport, getInsertPosition(insertParent, newImport, priority));
         PsiElement whitespace = element.getNextSibling();
         if (!(whitespace instanceof PsiWhiteSpace)) {
           whitespace = PsiParserFacade.SERVICE.getInstance(file.getProject()).createWhiteSpaceFromText("  >>> ");
@@ -350,10 +369,10 @@ public class AddImportHelper {
       }
       else {
         if (anchor instanceof PyImportStatementBase) {
-          insertParent.addAfter(nodeToInsert, anchor);
+          insertParent.addAfter(newImport, anchor);
         }
         else {
-          insertParent.addBefore(nodeToInsert, getInsertPosition(insertParent, nodeToInsert, priority));
+          insertParent.addBefore(newImport, getInsertPosition(insertParent, newImport, priority));
         }
       }
     }

diff --git a/python/src/com/jetbrains/python/inspections/quickfix/MoveFromFutureImportQuickFix.java b/python/src/com/jetbrains/python/inspections/quickfix/MoveFromFutureImportQuickFix.java
index 1fe8b3cb88d9714feefce39e4c78f2c638ee442e..d88008b081631c7538b6da5afde37e3a3a8f36e4 100644 (file)
--- a/python/src/com/jetbrains/python/inspections/quickfix/MoveFromFutureImportQuickFix.java
+++ b/python/src/com/jetbrains/python/inspections/quickfix/MoveFromFutureImportQuickFix.java
@@ -22,39 +22,32 @@ import com.intellij.psi.PsiElement;
 import com.intellij.psi.PsiFile;
 import com.jetbrains.python.PyBundle;
 import com.jetbrains.python.codeInsight.imports.AddImportHelper;
-import com.jetbrains.python.documentation.docstrings.DocStringUtil;
-import com.jetbrains.python.formatter.PyBlock;
 import com.jetbrains.python.psi.PyFile;
-import com.jetbrains.python.psi.PyStringLiteralExpression;
+import com.jetbrains.python.psi.PyFromImportStatement;
 import org.jetbrains.annotations.NotNull;
 
 /**
  * @author Alexey.Ivanov
  */
 public class MoveFromFutureImportQuickFix implements LocalQuickFix {
+  @Override
   @NotNull
   public String getName() {
     return PyBundle.message("QFIX.move.from.future.import");
   }
 
+  @Override
   @NotNull
   public String getFamilyName() {
     return getName();
   }
 
+  @Override
   public void applyFix(@NotNull Project project, @NotNull ProblemDescriptor descriptor) {
-    PsiElement problemElement = descriptor.getPsiElement();
-    PsiFile psiFile = problemElement.getContainingFile();
+    final PsiElement problemElement = descriptor.getPsiElement();
+    final PsiFile psiFile = problemElement.getContainingFile();
     if (psiFile instanceof PyFile) {
-      problemElement.putCopyableUserData(PyBlock.IMPORT_GROUP_BEGIN, true);
-      PyFile file = (PyFile)psiFile;
-      PyStringLiteralExpression docString = DocStringUtil.findDocStringExpression(file);
-      if (docString != null) {
-        file.addAfter(problemElement, docString.getParent() /* PyExpressionStatement */);
-      }
-      else {
-        file.addBefore(problemElement, file.getStatements().get(0));
-      }
+      AddImportHelper.addFromImportStatement(psiFile, (PyFromImportStatement)problemElement, AddImportHelper.ImportPriority.FUTURE, null);
       problemElement.delete();
     }
   }

diff --git a/python/testData/inspections/MoveFromFutureImportDocString_after.py b/python/testData/inspections/MoveFromFutureImportDocString_after.py
index 5f8f6e87ee0c6a6e0c90a8f34eaf0f8d80349236..b6fa176a74c3f2b9cc917a5747c680f739c5cf29 100644 (file)
--- a/python/testData/inspections/MoveFromFutureImportDocString_after.py
+++ b/python/testData/inspections/MoveFromFutureImportDocString_after.py
@@ -1,8 +1,10 @@
 """This is a docstring."""
-from __future__ import with_statement
+
 from __future__ import print_function
 #comment
 from __future__ import absolute_import
+from __future__ import with_statement
+
 
 class A:
     pass

diff --git a/python/testData/inspections/MoveFromFutureImport_after.py b/python/testData/inspections/MoveFromFutureImport_after.py
index b33e4a40b9a397b7ea8f18cde5ade434ec4e89af..b0f7e578abe3a950c0e805827d97a145a42419cd 100644 (file)
--- a/python/testData/inspections/MoveFromFutureImport_after.py
+++ b/python/testData/inspections/MoveFromFutureImport_after.py
@@ -1,7 +1,8 @@
-from __future__ import with_statement
 from __future__ import print_function
 #comment
 from __future__ import absolute_import
+from __future__ import with_statement
+
 
 class A:
     pass