do not propose unwrap for if in elif branch

fixed PY-8582 Unwrap for else doesn't propose unwrap for if, while for elif it does

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<PsiElement> 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);
   }

python/testData/refactoring/unwrap/ifInElifBranchUnwrap_before.py

@@ -0,0 +1,6 @@
+if "1":
+  print 1
+elif "2":
+  print<caret> 2
+else:
+  print 3

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<AnAction> 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());
+  }
 
 }