Constants are placed after docstrings in introduce constant refactoring (PY-3657)

python/src/com/jetbrains/python/codeInsight/imports/AddImportHelper.java

@@ -34,7 +34,11 @@ public class AddImportHelper {
     BUILTIN, THIRD_PARTY, PROJECT
   }
 
-  private static PsiElement getInsertPosition(final PsiFile file, String nameToImport, ImportPriority priority) {
+  public static PsiElement getFileInsertPosition(final PsiFile file) {
+    return getInsertPosition(file, null, null);
+  }
+
+  private static PsiElement getInsertPosition(final PsiFile file, @Nullable String nameToImport, @Nullable ImportPriority priority) {
     PsiElement feeler = file.getFirstChild();
     LOG.assertTrue(feeler != null);
     // skip initial comments and whitespace and try to get just below the last import stmt
@@ -43,7 +47,7 @@ public class AddImportHelper {
     PsiElement seeker = feeler;
     do {
       if (feeler instanceof PyImportStatementBase) {
-        if (shouldInsertBefore(file, (PyImportStatementBase)feeler, nameToImport, priority)) {
+        if (nameToImport != null && priority != null && shouldInsertBefore(file, (PyImportStatementBase)feeler, nameToImport, priority)) {
           break;
         }
         seeker = feeler;

python/src/com/jetbrains/python/refactoring/introduce/constant/ConstantIntroduceHandler.java

@@ -3,7 +3,9 @@ package com.jetbrains.python.refactoring.introduce.constant;
 import com.intellij.psi.PsiElement;
 import com.intellij.util.containers.HashSet;
 import com.jetbrains.python.PyBundle;
-import com.jetbrains.python.psi.*;
+import com.jetbrains.python.codeInsight.imports.AddImportHelper;
+import com.jetbrains.python.psi.PyExpression;
+import com.jetbrains.python.psi.PyFile;
 import com.jetbrains.python.refactoring.introduce.IntroduceHandler;
 import org.jetbrains.annotations.NotNull;
 
@@ -24,18 +26,9 @@ public class ConstantIntroduceHandler extends IntroduceHandler {
                                       @NotNull final List<PsiElement> occurrences,
                                       final boolean replaceAll,
                                       InitPlace initInConstructor) {
-    PsiElement anchor;
-    anchor = expression.getContainingFile();
+    final PsiElement anchor = expression.getContainingFile();
     assert anchor instanceof PyFile;
-    final List<PyStatement> statements = ((PyFile)anchor).getStatements();
-    int targetIndex = 0;
-    while(targetIndex < statements.size() && statements.get(targetIndex) instanceof PyImportStatementBase) {
-      targetIndex++;
-    }
-    if (targetIndex == statements.size()) {
-      return anchor.add(declaration);
-    }
-    return anchor.addBefore(declaration, statements.get(targetIndex));
+    return anchor.addBefore(declaration, AddImportHelper.getFileInsertPosition((PyFile)anchor));
   }
 
   public Collection<String> getSuggestedNames(@NotNull final PyExpression expression) {

python/testData/refactoring/introduceConstant/insertAfterDocstring.after.py

@@ -0,0 +1,8 @@
+"""
+Module docstring.
+"""
+a = 'Hello, World!'
+
+def f():
+    """Function docstring."""
+    print(a)

python/testData/refactoring/introduceConstant/insertAfterDocstring.py

@@ -0,0 +1,7 @@
+"""
+Module docstring.
+"""
+
+def f():
+    """Function docstring."""
+    print('<caret>Hello, World!')

python/testSrc/com/jetbrains/python/refactoring/PyIntroduceConstantTest.java

@@ -26,6 +26,10 @@ public class PyIntroduceConstantTest extends PyLightFixtureTestCase {
     doTest();
   }
 
+  public void testInsertAfterDocstring() { // PY-3657
+    doTest();
+  }
+
   private void doTest() {
     myFixture.configureByFile("/refactoring/introduceConstant/" + getTestName(true) + ".py");
     ConstantIntroduceHandler handler = new ConstantIntroduceHandler();