Make "extract method" work on multi-line function calls

Author: chbndrhnns

python/python-psi-impl/src/com/jetbrains/python/refactoring/PyRefactoringUtil.java

@@ -75,6 +75,15 @@ public void visitElement(final @NotNull PsiElement element) {
     if (parent == null) {
       return null;
     }
+
+    // Allow function call where arguments are on multiple lines
+    if (parent instanceof PyArgumentList) {
+      PyCallExpression callExpression = PsiTreeUtil.getParentOfType(parent, PyCallExpression.class);
+      if (callExpression != null) {
+        return callExpression;
+      }
+    }
+
     // If it is PyIfPart for example, parent if statement, we should deny
     if (!(parent instanceof PyExpression)){
       return null;

python/python-psi-impl/src/com/jetbrains/python/refactoring/extractmethod/PyExtractMethodHandler.java

@@ -14,11 +14,13 @@
 import com.intellij.refactoring.RefactoringBundle;
 import com.intellij.refactoring.util.CommonRefactoringUtil;
 import com.jetbrains.python.PyPsiBundle;
+import com.jetbrains.python.PyTokenTypes;
 import com.jetbrains.python.codeInsight.codeFragment.PyCodeFragment;
 import com.jetbrains.python.codeInsight.codeFragment.PyCodeFragmentUtil;
 import com.jetbrains.python.codeInsight.controlflow.ScopeOwner;
 import com.jetbrains.python.psi.PyClass;
 import com.jetbrains.python.psi.PyElement;
+import com.jetbrains.python.psi.PyElementType;
 import com.jetbrains.python.psi.impl.PyPsiUtils;
 import com.jetbrains.python.refactoring.PyRefactoringUtil;
 import org.jetbrains.annotations.NotNull;
@@ -143,6 +145,17 @@ private static boolean rangeBelongsToSameClassBody(@NotNull PsiElement element1,
         element2 == PyPsiUtils.getPrevSignificantLeaf(PsiTreeUtil.getDeepestLast(statement2), !(element2 instanceof PsiComment))) {
       return Couple.of(statement1, statement2);
     }
+
+    // multi-line function call
+    boolean isLeftParen = "(".equals(element2.getText());
+    PsiElement prevSignificantLeaf = PyPsiUtils.getPrevSignificantLeaf(
+      PsiTreeUtil.getDeepestLast(statement2),
+      !(statement2 instanceof PsiComment)
+    );
+    boolean isRightParen = prevSignificantLeaf != null && ")".equals(prevSignificantLeaf.getText());
+    if (isLeftParen && isRightParen) {
+      return Couple.of(statement1, statement2);
+    }
     return null;
   }
 }

python/testData/refactoring/extractmethod/StartInArgumentListOfMultiLineFunctionCall.after.py

@@ -0,0 +1,14 @@
+def long_function_name(**kwargs): ...
+
+def example_function():
+    result = extracted()
+    processed = result.upper()
+    return processed
+
+
+def extracted():
+    return long_function_name(
+        first_argument="value1",
+        second_argument="value2",
+        third_argument="value3"
+    )

python/testData/refactoring/extractmethod/StartInArgumentListOfMultiLineFunctionCall.before.py

@@ -0,0 +1,10 @@
+def long_function_name(**kwargs): ...
+
+def example_function():
+    result = long_function_name(
+        first_argument<caret>="value1",
+        second_argument="value2",
+        third_argument="value3"
+    )
+    processed = result.upper()
+    return processed

python/testData/refactoring/extractmethod/StartOnMultiLineFunctionCall.after.py

@@ -0,0 +1,15 @@
+def long_function_name(**kwargs): ...
+
+def example_function():
+    result = extracted()
+    processed = result.upper()
+    return processed
+
+
+def extracted():
+    result = long_function_name(
+        first_argument="value1",
+        second_argument="value2",
+        third_argument="value3"
+    )
+    return result

python/testData/refactoring/extractmethod/StartOnMultiLineFunctionCall.before.py

@@ -0,0 +1,10 @@
+def long_function_name(**kwargs): ...
+
+def example_function():
+    result = long_function_<caret>name(
+        first_argument="value1",
+        second_argument="value2",
+        third_argument="value3"
+    )
+    processed = result.upper()
+    return processed

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

@@ -68,6 +68,16 @@ public void testMethodInnerFuncCombined() {
     doTest("extracted");
   }
 
+  // PY-53711
+  public void testStartInArgumentListOfMultiLineFunctionCall() {
+    doTest("extracted");
+  }
+
+  // PY-53711
+  public void testStartOnMultiLineFunctionCall() {
+    doTest("extracted");
+  }
+
   public void testParameter() {
     doTest("bar");
   }
@@ -315,4 +325,4 @@ public void testCommentsPrecedingSourceStatement() {
   public void testInterruptedOuterLoop() {
     doFail("foo", "Cannot perform refactoring when execution flow is interrupted");
   }
-}
+}