diff --git a/python/psi-api/src/com/jetbrains/python/psi/impl/PyPsiUtils.java b/python/psi-api/src/com/jetbrains/python/psi/impl/PyPsiUtils.java index 8c2ea8771944..7a00de223517 100644 --- a/python/psi-api/src/com/jetbrains/python/psi/impl/PyPsiUtils.java +++ b/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() { + @Override + public boolean value(PsiElement element) { + return element.getParent() == superParent; + } + }); } public static List 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; } diff --git a/python/src/com/jetbrains/python/refactoring/extractmethod/PyExtractMethodHandler.java b/python/src/com/jetbrains/python/refactoring/extractmethod/PyExtractMethodHandler.java index cdad134dc530..0af213341578 100644 --- a/python/src/com/jetbrains/python/refactoring/extractmethod/PyExtractMethodHandler.java +++ b/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 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 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; } diff --git a/python/src/com/jetbrains/python/refactoring/extractmethod/PyExtractMethodUtil.java b/python/src/com/jetbrains/python/refactoring/extractmethod/PyExtractMethodUtil.java index cc73ab4869ee..59a2db24cb93 100644 --- a/python/src/com/jetbrains/python/refactoring/extractmethod/PyExtractMethodUtil.java +++ b/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());