PY-17265 Placement of inserted function is consistent with Move refactoring

Namely, if a function is being moved to another file it will be
inserted at its end unless there is a top-level usage of it. In this
case the generated function will be inserted right before the first
such usage so as not to produce unresolved references. If function
stays in the same file, it will be inserted after its original parent
statement i.e. another function or a class, again if there is no
conflicting usages.

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

@@ -151,6 +151,15 @@ public class AddImportHelper {
     return PsiTreeUtil.getParentOfType(anchor, PyStatement.class, false);
   }
 
+  /**
+   * Returns position in the file after all leading comments, docstring and import statements.
+   * <p>
+   * Returned PSI element is intended to be used as "anchor" parameter for {@link PsiElement#addBefore(PsiElement, PsiElement)},
+   * hence {@code null} means that element to be inserted will be the first in the file.
+   *
+   * @param file target file where some new top-level element is going to be inserted
+   * @return anchor PSI element as described
+   */
   @Nullable
   public static PsiElement getFileInsertPosition(final PsiFile file) {
     return getInsertPosition(file, null, null);

python/src/com/jetbrains/python/refactoring/move/PyMoveRefactoringUtil.java

@@ -18,15 +18,26 @@ package com.jetbrains.python.refactoring.move;
 import com.intellij.openapi.util.text.StringUtil;
 import com.intellij.openapi.vfs.VirtualFile;
 import com.intellij.psi.PsiElement;
+import com.intellij.psi.PsiFile;
 import com.intellij.psi.PsiNamedElement;
+import com.intellij.psi.util.PsiUtilCore;
 import com.intellij.psi.util.QualifiedName;
+import com.intellij.usageView.UsageInfo;
 import com.intellij.util.IncorrectOperationException;
 import com.jetbrains.python.PyBundle;
+import com.jetbrains.python.codeInsight.dataflow.scope.ScopeUtil;
 import com.jetbrains.python.psi.PyFunction;
 import com.jetbrains.python.psi.PyQualifiedNameOwner;
+import com.jetbrains.python.psi.impl.PyImportStatementNavigator;
+import com.jetbrains.python.psi.impl.PyPsiUtils;
 import com.jetbrains.python.psi.resolve.QualifiedNameFinder;
 import com.jetbrains.python.refactoring.classes.PyClassRefactoringUtil;
 import org.jetbrains.annotations.NotNull;
+import org.jetbrains.annotations.Nullable;
+
+import java.util.List;
+import java.util.Objects;
+import java.util.Optional;
 
 /**
  * @author Mikhail Golubev
@@ -78,4 +89,31 @@ public class PyMoveRefactoringUtil {
     }
     return "";
   }
+
+  /**
+   * Returns anchor PSI element for {@link PsiElement#addBefore(PsiElement, PsiElement)}.
+   * <p>
+   * If there are any usages at file's level returns the top-level parent element for the first of them,
+   * otherwise return {@code null} which means that the element can be safely inserted at the end of the file.
+   *
+   * @param usages      usages of the original element
+   * @param destination file where original/generated element is to be moved
+   * @return anchor element as described
+   */
+  @Nullable
+  public static PsiElement findLowestPossibleTopLevelInsertionPosition(@NotNull List<UsageInfo> usages, @NotNull PsiFile destination) {
+    return findFirstTopLevelUsageInFile(usages, destination)
+      .map(element -> PyPsiUtils.getParentRightBefore(element, element.getContainingFile()))
+      .orElse(null);
+  }
+
+  @NotNull
+  private static Optional<PsiElement> findFirstTopLevelUsageInFile(@NotNull List<UsageInfo> usages, @NotNull PsiFile destination) {
+    return usages.stream()
+      .map(UsageInfo::getElement)
+      .filter(Objects::nonNull)
+      .filter(element -> ScopeUtil.getScopeOwner(element) == destination)
+      .filter(element -> PyImportStatementNavigator.getImportStatementByElement(element) == null)
+      .min(PsiUtilCore::compareElementsByPosition);
+  }
 }

python/src/com/jetbrains/python/refactoring/move/makeFunctionTopLevel/PyBaseMakeFunctionTopLevelProcessor.java

@@ -37,7 +37,6 @@ import com.jetbrains.python.codeInsight.controlflow.ControlFlowCache;
 import com.jetbrains.python.codeInsight.controlflow.ReadWriteInstruction;
 import com.jetbrains.python.codeInsight.controlflow.ScopeOwner;
 import com.jetbrains.python.codeInsight.dataflow.scope.ScopeUtil;
-import com.jetbrains.python.codeInsight.imports.AddImportHelper;
 import com.jetbrains.python.psi.*;
 import com.jetbrains.python.psi.impl.PyPsiUtils;
 import com.jetbrains.python.psi.resolve.PyResolveContext;
@@ -46,11 +45,9 @@ import com.jetbrains.python.refactoring.PyRefactoringUtil;
 import com.jetbrains.python.refactoring.classes.PyClassRefactoringUtil;
 import com.jetbrains.python.refactoring.move.PyMoveRefactoringUtil;
 import org.jetbrains.annotations.NotNull;
+import org.jetbrains.annotations.Nullable;
 
-import java.util.ArrayList;
-import java.util.Collection;
-import java.util.List;
-import java.util.Set;
+import java.util.*;
 
 import static com.jetbrains.python.psi.PyUtil.as;
 
@@ -117,11 +114,12 @@ public abstract class PyBaseMakeFunctionTopLevelProcessor extends BaseRefactorin
     if (importsRequired(usages, targetFile)) {
       PyMoveRefactoringUtil.checkValidImportableFile(targetFile, targetFile.getVirtualFile());
     }
-    
+
+    final PsiElement position = PyMoveRefactoringUtil.findLowestPossibleTopLevelInsertionPosition(Arrays.asList(usages), targetFile);
     // We should update usages before we generate and insert new function, because we have to update its usages inside 
     // (e.g. recursive calls) it first 
     updateUsages(newParameters, usages);
-    final PyFunction newFunction = insertFunction(createNewFunction(newParameters), targetFile);
+    final PyFunction newFunction = insertFunction(createNewFunction(newParameters), targetFile, position);
 
     myFunction.delete();
 
@@ -200,17 +198,16 @@ public abstract class PyBaseMakeFunctionTopLevelProcessor extends BaseRefactorin
   }
 
   @NotNull
-  protected PyFunction insertFunction(@NotNull PyFunction newFunction, PyFile newFile) {
-    final PyFunction replacement;
+  protected PyFunction insertFunction(@NotNull PyFunction newFunction, @NotNull PyFile newFile, @Nullable PsiElement anchor) {
     if (mySourceFile == newFile) {
-      final PsiElement anchor;
-      anchor = PyPsiUtils.getParentRightBefore(myFunction, mySourceFile);
-      replacement = (PyFunction)mySourceFile.addAfter(newFunction, anchor);
+      // In the same file try inserting generated function at the top level but preferably right after the original scope owner
+      final PsiElement surroundingStatement = PyPsiUtils.getParentRightBefore(myFunction, mySourceFile);
+      if (anchor == null || surroundingStatement.getTextRange().getEndOffset() < anchor.getTextRange().getStartOffset()) {
+        return (PyFunction)mySourceFile.addAfter(newFunction, surroundingStatement);
+      }
     }
-    else {
-      replacement = (PyFunction)newFile.addAfter(newFunction, AddImportHelper.getFileInsertPosition(newFile));
-    }
-    return replacement;
+    // Insert at the end or before first top-level usage in the file
+    return (PyFunction)newFile.addBefore(newFunction, anchor);
   }
 
   @NotNull

python/src/com/jetbrains/python/refactoring/move/moduleMembers/PyMoveSymbolProcessor.java

@@ -29,6 +29,7 @@ import com.jetbrains.python.psi.*;
 import com.jetbrains.python.psi.resolve.PyResolveContext;
 import com.jetbrains.python.psi.resolve.QualifiedNameFinder;
 import com.jetbrains.python.refactoring.classes.PyClassRefactoringUtil;
+import com.jetbrains.python.refactoring.move.PyMoveRefactoringUtil;
 import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;
 
@@ -94,33 +95,8 @@ public class PyMoveSymbolProcessor {
 
   @NotNull
   private PsiElement addElementToFile(@NotNull PsiElement element) {
-    final PsiElement firstUsage = findFirstTopLevelWithUsageAtDestination();
-    if (firstUsage != null) {
-      return myDestinationFile.addBefore(element, firstUsage);
-    }
-    else {
-      return myDestinationFile.add(element);
-    }
-  }
-
-  @Nullable
-  private PsiElement findFirstTopLevelWithUsageAtDestination() {
-    final List<PsiElement> topLevelAtDestination = ContainerUtil.mapNotNull(myUsages, usage -> {
-      final PsiElement element = usage.getElement();
-      if (element != null && ScopeUtil.getScopeOwner(element) == myDestinationFile && getImportStatementByElement(element) == null) {
-        return findTopLevelParent(element);
-      }
-      return null;
-    });
-    if (topLevelAtDestination.isEmpty()) {
-      return null;
-    }
-    return Collections.min(topLevelAtDestination, PsiUtilCore::compareElementsByPosition);
-  }
-
-  @Nullable
-  private PsiElement findTopLevelParent(@NotNull PsiElement element) {
-    return PsiTreeUtil.findFirstParent(element, element1 -> element1.getParent() == myDestinationFile);
+    final PsiElement anchor = PyMoveRefactoringUtil.findLowestPossibleTopLevelInsertionPosition(myUsages, myDestinationFile);
+    return myDestinationFile.addBefore(element, anchor);
   }
 
   private void updateSingleUsage(@NotNull PsiElement usage, @NotNull PsiNamedElement newElement) {

python/testData/refactoring/makeFunctionTopLevel/methodInsertionPositionNoUsageInAnotherFile/after/main.py

@@ -0,0 +1,2 @@
+class C:
+    pass

python/testData/refactoring/makeFunctionTopLevel/methodInsertionPositionNoUsageInAnotherFile/after/other.py

@@ -0,0 +1,14 @@
+def already_existing1():
+    pass
+
+
+def already_existing2():
+    pass
+
+
+def already_existing3():
+    pass
+
+
+def method():
+    pass

python/testData/refactoring/makeFunctionTopLevel/methodInsertionPositionNoUsageInAnotherFile/before/main.py

@@ -0,0 +1,3 @@
+class C:
+    def met<caret>hod(self):
+        pass

python/testData/refactoring/makeFunctionTopLevel/methodInsertionPositionNoUsageInAnotherFile/before/other.py

@@ -0,0 +1,10 @@
+def already_existing1():
+    pass
+
+
+def already_existing2():
+    pass
+
+
+def already_existing3():
+    pass

python/testData/refactoring/makeFunctionTopLevel/methodInsertionPositionSameFileClassAndUsageNotTopLevel.after.py

@@ -0,0 +1,10 @@
+def method():
+    pass
+
+
+if True:
+    class C:
+        pass
+
+
+    method()

python/testData/refactoring/makeFunctionTopLevel/methodInsertionPositionSameFileClassAndUsageNotTopLevel.py

@@ -0,0 +1,7 @@
+if True:
+    class C:
+        def met<caret>hod(self):
+            pass
+
+
+    C().method()

python/testData/refactoring/makeFunctionTopLevel/methodInsertionPositionUsageInAnotherFile/after/main.py

@@ -0,0 +1,2 @@
+class C:
+    pass

python/testData/refactoring/makeFunctionTopLevel/methodInsertionPositionUsageInAnotherFile/after/other.py

@@ -0,0 +1,20 @@
+def already_existing1():
+    pass
+
+
+def already_existing2():
+    pass
+
+
+from main import C
+
+
+def method():
+    pass
+
+
+method()
+
+
+def already_existing3():
+    pass

python/testData/refactoring/makeFunctionTopLevel/methodInsertionPositionUsageInAnotherFile/before/main.py

@@ -0,0 +1,3 @@
+class C:
+    def met<caret>hod(self):
+        pass

python/testData/refactoring/makeFunctionTopLevel/methodInsertionPositionUsageInAnotherFile/before/other.py

@@ -0,0 +1,15 @@
+def already_existing1():
+    pass
+
+
+def already_existing2():
+    pass
+
+
+from main import C
+
+C().method()
+
+
+def already_existing3():
+    pass

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

@@ -23,11 +23,7 @@ import com.intellij.openapi.vfs.VirtualFile;
 import com.intellij.psi.PsiElement;
 import com.intellij.psi.PsiFile;
 import com.intellij.psi.PsiReference;
-import com.intellij.refactoring.RefactoringBundle;
-import com.intellij.refactoring.actions.MoveAction;
-import com.intellij.refactoring.util.CommonRefactoringUtil;
 import com.intellij.testFramework.PlatformTestUtil;
-import com.intellij.testFramework.TestActionEvent;
 import com.intellij.util.IncorrectOperationException;
 import com.jetbrains.python.PyBundle;
 import com.jetbrains.python.PyTokenTypes;
@@ -276,6 +272,18 @@ public class PyMakeFunctionTopLevelTest extends PyTestCase {
     doTestFailure(PyBundle.message("refactoring.move.error.destination.file.contains.function.$0", "nested"));
   }
 
+  public void testMethodInsertionPositionSameFileClassAndUsageNotTopLevel() {
+    doTestSuccess();
+  }
+
+  public void testMethodInsertionPositionUsageInAnotherFile() throws IOException {
+    doMultiFileTest("other.py", null);
+  }
+
+  public void testMethodInsertionPositionNoUsageInAnotherFile() throws IOException {
+    doMultiFileTest("other.py", null);    
+  }
+
   @Override
   protected String getTestDataPath() {
     return super.getTestDataPath() + "/refactoring/makeFunctionTopLevel/";