diff --git a/python/src/com/jetbrains/python/refactoring/unwrap/PyIfUnwrapper.java b/python/src/com/jetbrains/python/refactoring/unwrap/PyIfUnwrapper.java index 7d591d74bcee..9a546a7c833a 100644 --- a/python/src/com/jetbrains/python/refactoring/unwrap/PyIfUnwrapper.java +++ b/python/src/com/jetbrains/python/refactoring/unwrap/PyIfUnwrapper.java @@ -1,12 +1,13 @@ package com.jetbrains.python.refactoring.unwrap; import com.intellij.psi.PsiElement; +import com.intellij.psi.util.PsiTreeUtil; import com.intellij.util.IncorrectOperationException; import com.jetbrains.python.PyBundle; -import com.jetbrains.python.psi.PyIfStatement; -import com.jetbrains.python.psi.PyPassStatement; -import com.jetbrains.python.psi.PyStatement; -import com.jetbrains.python.psi.PyStatementList; +import com.jetbrains.python.psi.*; +import com.jetbrains.python.psi.impl.PyIfPartIfImpl; + +import java.util.List; /** * User : ktisha @@ -17,8 +18,8 @@ public class PyIfUnwrapper extends PyUnwrapper { } public boolean isApplicableTo(PsiElement e) { - if (e instanceof PyIfStatement) { - final PyStatementList statementList = ((PyIfStatement)e).getIfPart().getStatementList(); + if (e instanceof PyIfPartIfImpl) { + final PyStatementList statementList = ((PyIfPartIfImpl)e).getStatementList(); if (statementList != null) { final PyStatement[] statements = statementList.getStatements(); return statements.length == 1 && !(statements[0] instanceof PyPassStatement) || statements.length > 1; @@ -27,9 +28,16 @@ public class PyIfUnwrapper extends PyUnwrapper { return false; } + @Override + public PsiElement collectAffectedElements(PsiElement e, List toExtract) { + super.collectAffectedElements(e, toExtract); + return PsiTreeUtil.getParentOfType(e, PyIfStatement.class); + } + + @Override protected void doUnwrap(final PsiElement element, final Context context) throws IncorrectOperationException { - final PyIfStatement ifStatement = (PyIfStatement)element; + final PyIfStatement ifStatement = PsiTreeUtil.getParentOfType(element, PyIfStatement.class); context.extractPart(ifStatement); context.delete(ifStatement); } diff --git a/python/testData/refactoring/unwrap/ifInElifBranchUnwrap_before.py b/python/testData/refactoring/unwrap/ifInElifBranchUnwrap_before.py new file mode 100644 index 000000000000..995434072cfc --- /dev/null +++ b/python/testData/refactoring/unwrap/ifInElifBranchUnwrap_before.py @@ -0,0 +1,6 @@ +if "1": + print 1 +elif "2": + print 2 +else: + print 3 \ No newline at end of file diff --git a/python/testSrc/com/jetbrains/python/refactoring/PyUnwrapperTest.java b/python/testSrc/com/jetbrains/python/refactoring/PyUnwrapperTest.java index d35468ce7445..17d5113e1873 100644 --- a/python/testSrc/com/jetbrains/python/refactoring/PyUnwrapperTest.java +++ b/python/testSrc/com/jetbrains/python/refactoring/PyUnwrapperTest.java @@ -4,6 +4,7 @@ import com.intellij.codeInsight.unwrap.UnwrapHandler; import com.intellij.openapi.actionSystem.AnAction; import com.intellij.openapi.editor.Editor; import com.intellij.psi.PsiFile; +import com.jetbrains.python.PyBundle; import com.jetbrains.python.fixtures.PyTestCase; import com.jetbrains.python.psi.LanguageLevel; @@ -44,6 +45,8 @@ public class PyUnwrapperTest extends PyTestCase { public void testEndOfStatementUnwrap() throws Throwable {doTest();} public void testEndOfStatementNextLineUnwrap() throws Throwable {doNegativeTest();} + public void testIfInElifBranchUnwrap() throws Throwable {doNegativeTest(PyBundle.message("unwrap.if"));} + private void doTest() { doTest(0); } @@ -85,6 +88,20 @@ public class PyUnwrapperTest extends PyTestCase { } }; h.invoke(myFixture.getProject(), myFixture.getEditor(), myFixture.getFile()); - } + } + + private void doNegativeTest(final String optionName) { + String before = "refactoring/unwrap/" + getTestName(true) + "_before.py"; + myFixture.configureByFile(before); + UnwrapHandler h = new UnwrapHandler() { + @Override + protected void selectOption(List options, Editor editor, PsiFile file) { + for (AnAction option : options) { + assertFalse("\"" + optionName + "\" is available to unwrap ", option.toString().contains(optionName)); + } + } + }; + h.invoke(myFixture.getProject(), myFixture.getEditor(), myFixture.getFile()); + } }