diff --git a/python/src/com/jetbrains/python/codeInsight/imports/AddImportHelper.java b/python/src/com/jetbrains/python/codeInsight/imports/AddImportHelper.java index 41d91bf59b61..f4e1c7e97b04 100644 --- a/python/src/com/jetbrains/python/codeInsight/imports/AddImportHelper.java +++ b/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; diff --git a/python/src/com/jetbrains/python/refactoring/introduce/constant/ConstantIntroduceHandler.java b/python/src/com/jetbrains/python/refactoring/introduce/constant/ConstantIntroduceHandler.java index d3e981daa7d9..7f4483180548 100644 --- a/python/src/com/jetbrains/python/refactoring/introduce/constant/ConstantIntroduceHandler.java +++ b/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 occurrences, final boolean replaceAll, InitPlace initInConstructor) { - PsiElement anchor; - anchor = expression.getContainingFile(); + final PsiElement anchor = expression.getContainingFile(); assert anchor instanceof PyFile; - final List 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 getSuggestedNames(@NotNull final PyExpression expression) { diff --git a/python/testData/refactoring/introduceConstant/insertAfterDocstring.after.py b/python/testData/refactoring/introduceConstant/insertAfterDocstring.after.py new file mode 100644 index 000000000000..145fb3e8667c --- /dev/null +++ b/python/testData/refactoring/introduceConstant/insertAfterDocstring.after.py @@ -0,0 +1,8 @@ +""" +Module docstring. +""" +a = 'Hello, World!' + +def f(): + """Function docstring.""" + print(a) \ No newline at end of file diff --git a/python/testData/refactoring/introduceConstant/insertAfterDocstring.py b/python/testData/refactoring/introduceConstant/insertAfterDocstring.py new file mode 100644 index 000000000000..1b1c15964c48 --- /dev/null +++ b/python/testData/refactoring/introduceConstant/insertAfterDocstring.py @@ -0,0 +1,7 @@ +""" +Module docstring. +""" + +def f(): + """Function docstring.""" + print('Hello, World!') \ No newline at end of file diff --git a/python/testSrc/com/jetbrains/python/refactoring/PyIntroduceConstantTest.java b/python/testSrc/com/jetbrains/python/refactoring/PyIntroduceConstantTest.java index 5fdac50e796e..c0ccdfbdd217 100644 --- a/python/testSrc/com/jetbrains/python/refactoring/PyIntroduceConstantTest.java +++ b/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();