Slightly refactor PyExtractMethodHandler

* Renamed PyPsiUtil#getStatement to PyPsiUtil#getParentRightBefore as
this name more clearly describes method's functionality.
* Return Couple<PsiElement> instead of awkward two-items array or
PsiElements.

python/psi-api/src/com/jetbrains/python/psi/impl/PyPsiUtils.java

@@ -17,6 +17,7 @@ package com.jetbrains.python.psi.impl;
 
 import com.intellij.lang.ASTNode;
 import com.intellij.openapi.diagnostic.Logger;
+import com.intellij.openapi.util.Condition;
 import com.intellij.openapi.util.text.StringUtil;
 import com.intellij.psi.*;
 import com.intellij.psi.stubs.StubElement;
@@ -115,7 +116,7 @@ public class PyPsiUtils {
     if (compStatement == null) {
       return null;
     }
-    return getStatement(compStatement, element);
+    return getParentRightBefore(element, compStatement);
   }
 
   public static PyElement getStatementList(final PsiElement element) {
@@ -125,14 +126,21 @@ public class PyPsiUtils {
            : PsiTreeUtil.getParentOfType(element, PyFile.class, PyStatementList.class);
   }
 
+  /**
+   * Returns ancestor of the element that is also direct child of the given super parent.
+   *
+   * @param element     element to start search from
+   * @param superParent direct parent of the desired ancestor
+   * @return described element or {@code null} if it doesn't exist
+   */
   @Nullable
-  public static PsiElement getStatement(final PsiElement compStatement, PsiElement element) {
-    PsiElement parent = element.getParent();
-    while (parent != null && parent != compStatement) {
-      element = parent;
-      parent = element.getParent();
-    }
-    return parent != null ? element : null;
+  public static PsiElement getParentRightBefore(@NotNull PsiElement element, @NotNull final PsiElement superParent) {
+    return PsiTreeUtil.findFirstParent(element, false, new Condition<PsiElement>() {
+      @Override
+      public boolean value(PsiElement element) {
+        return element.getParent() == superParent;
+      }
+    });
   }
 
   public static List<PsiElement> collectElements(final PsiElement statement1, final PsiElement statement2) {
@@ -161,7 +169,7 @@ public class PyPsiUtils {
 
   public static int getElementIndentation(final PsiElement element) {
     final PsiElement compStatement = getStatementList(element);
-    final PsiElement statement = getStatement(compStatement, element);
+    final PsiElement statement = getParentRightBefore(element, compStatement);
     if (statement == null) {
       return 0;
     }

python/src/com/jetbrains/python/refactoring/extractmethod/PyExtractMethodHandler.java

@@ -19,6 +19,7 @@ import com.intellij.codeInsight.codeFragment.CannotCreateCodeFragmentException;
 import com.intellij.openapi.actionSystem.DataContext;
 import com.intellij.openapi.editor.*;
 import com.intellij.openapi.project.Project;
+import com.intellij.openapi.util.Couple;
 import com.intellij.psi.PsiComment;
 import com.intellij.psi.PsiElement;
 import com.intellij.psi.PsiFile;
@@ -82,9 +83,9 @@ public class PyExtractMethodHandler implements RefactoringActionHandler {
       return;
     }
 
-    final PsiElement[] statements = getStatementsRange(element1, element2);
+    final Couple<PsiElement> statements = getStatementsRange(element1, element2);
     if (statements != null) {
-      final ScopeOwner owner = PsiTreeUtil.getParentOfType(statements[0], ScopeOwner.class);
+      final ScopeOwner owner = PsiTreeUtil.getParentOfType(statements.getFirst(), ScopeOwner.class);
       if (owner == null) {
         return;
       }
@@ -97,7 +98,7 @@ public class PyExtractMethodHandler implements RefactoringActionHandler {
                                             RefactoringBundle.message("extract.method.title"), "refactoring.extractMethod");
         return;
       }
-      PyExtractMethodUtil.extractFromStatements(project, editor, fragment, statements[0], statements[1]);
+      PyExtractMethodUtil.extractFromStatements(project, editor, fragment, statements.getFirst(), statements.getSecond());
       return;
     }
 
@@ -126,7 +127,7 @@ public class PyExtractMethodHandler implements RefactoringActionHandler {
   }
 
   @Nullable
-  private static PsiElement[] getStatementsRange(final PsiElement element1, final PsiElement element2) {
+  private static Couple<PsiElement> getStatementsRange(final PsiElement element1, final PsiElement element2) {
     final PsiElement parent = PsiTreeUtil.findCommonParent(element1, element2);
     if (parent == null) {
       return null;
@@ -137,8 +138,8 @@ public class PyExtractMethodHandler implements RefactoringActionHandler {
       return null;
     }
 
-    final PsiElement statement1 = PyPsiUtils.getStatement(statementList, element1);
-    final PsiElement statement2 = PyPsiUtils.getStatement(statementList, element2);
+    final PsiElement statement1 = PyPsiUtils.getParentRightBefore(element1, statementList);
+    final PsiElement statement2 = PyPsiUtils.getParentRightBefore(element2, statementList);
     if (statement1 == null || statement2 == null){
       return null;
     }
@@ -146,7 +147,7 @@ public class PyExtractMethodHandler implements RefactoringActionHandler {
     // return elements if they are really first and last elements of statements
     if (element1 == PsiTreeUtil.getDeepestFirst(statement1) &&
         element2 == PyPsiUtils.getSignificantToTheLeft(PsiTreeUtil.getDeepestLast(statement2), !(element2 instanceof PsiComment))) {
-      return new PsiElement[]{statement1, statement2};
+      return Couple.of(statement1, statement2);
     }
     return null;
   }

python/src/com/jetbrains/python/refactoring/extractmethod/PyExtractMethodUtil.java

@@ -479,7 +479,7 @@ public class PyExtractMethodUtil {
     PsiElement result;
     if (parent instanceof PyFile || parent instanceof PyClass) {
       PsiElement target = parent instanceof PyClass ? ((PyClass)parent).getStatementList() : parent;
-      final PsiElement anchorStatement = PyPsiUtils.getStatement(target, anchor);
+      final PsiElement anchorStatement = PyPsiUtils.getParentRightBefore(anchor, target);
       result = target.addBefore(generatedMethod, anchorStatement);
     }
     else {
@@ -635,7 +635,7 @@ public class PyExtractMethodUtil {
 
     //return if don`t want to extract method
     if (!dialog.isOK()) {
-      return Pair.create(null, null);
+      return Pair.empty();
     }
 
     return Pair.create(dialog.getMethodName(), dialog.getVariableData());