fixed PY-8566 Unwrap/remove action is not available for with statement

python/src/com/jetbrains/python/PyBundle.properties

@@ -222,6 +222,7 @@ surround.with.try.except.template=try / except
 unwrap.if=Unwrap if...
 unwrap.while=Unwrap while...
 unwrap.for=Unwrap for...
+unwrap.with=Unwrap with...
 unwrap.try=Unwrap try...
 unwrap.else=Unwrap else...
 unwrap.elif=Unwrap elif...

python/src/com/jetbrains/python/refactoring/unwrap/PyUnwrapDescriptor.java

@@ -17,7 +17,8 @@ public class PyUnwrapDescriptor extends UnwrapDescriptorBase{
       new PyElIfUnwrapper(),
       new PyElIfRemover(),
       new PyTryUnwrapper(),
-      new PyForUnwrapper()
+      new PyForUnwrapper(),
+      new PyWithUnwrapper()
     };
   }
 }

python/src/com/jetbrains/python/refactoring/unwrap/PyUnwrapper.java

@@ -1,11 +1,13 @@
 package com.jetbrains.python.refactoring.unwrap;
 
 import com.intellij.codeInsight.unwrap.AbstractUnwrapper;
+import com.intellij.lang.ASTNode;
 import com.intellij.openapi.editor.Editor;
 import com.intellij.psi.PsiElement;
 import com.intellij.psi.PsiWhiteSpace;
 import com.intellij.psi.impl.source.codeStyle.CodeEditUtil;
 import com.intellij.util.IncorrectOperationException;
+import com.jetbrains.python.PyElementTypes;
 import com.jetbrains.python.psi.*;
 import com.jetbrains.python.psi.impl.PyIfPartElifImpl;
 import com.jetbrains.python.psi.impl.PyIfPartIfImpl;
@@ -47,6 +49,9 @@ public abstract class PyUnwrapper extends AbstractUnwrapper<PyUnwrapper.Context>
       else if (from instanceof PyIfPartElifImpl) {
         extractFromElseBlock((PyIfPartElifImpl)from);
       }
+      else if (from instanceof PyWithStatement) {
+        extractFromWithBlock((PyWithStatement)from);
+      }
     }
 
     public void extractFromConditionalBlock(PyStatementWithElse from) {
@@ -79,6 +84,15 @@ public abstract class PyUnwrapper extends AbstractUnwrapper<PyUnwrapper.Context>
         extract(body.getFirstChild(), body.getLastChild(), from.getParent());
     }
 
+    public void extractFromWithBlock(PyWithStatement from) {
+      ASTNode n = from.getNode().findChildByType(PyElementTypes.STATEMENT_LISTS);
+      if (n != null) {
+        final PyStatementList body = (PyStatementList)n.getPsi();
+        if (body != null)
+          extract(body.getFirstChild(), body.getLastChild(), from);
+      }
+    }
+
     @Override
     protected boolean isWhiteSpace(PsiElement element) {
       return element instanceof PsiWhiteSpace;

python/src/com/jetbrains/python/refactoring/unwrap/PyWithUnwrapper.java

@@ -0,0 +1,39 @@
+package com.jetbrains.python.refactoring.unwrap;
+
+import com.intellij.lang.ASTNode;
+import com.intellij.psi.PsiElement;
+import com.intellij.util.IncorrectOperationException;
+import com.jetbrains.python.PyBundle;
+import com.jetbrains.python.PyElementTypes;
+import com.jetbrains.python.psi.*;
+
+/**
+ * User : ktisha
+ */
+public class PyWithUnwrapper extends PyUnwrapper {
+  public PyWithUnwrapper() {
+    super(PyBundle.message("unwrap.with"));
+  }
+
+  public boolean isApplicableTo(PsiElement e) {
+    if (e instanceof PyWithStatement) {
+      ASTNode n = e.getNode().findChildByType(PyElementTypes.STATEMENT_LISTS);
+      if (n != null) {
+        final PyStatementList statementList = (PyStatementList)n.getPsi();
+        if (statementList != null) {
+          final PyStatement[] statements = statementList.getStatements();
+          return statements.length == 1 && !(statements[0] instanceof PyPassStatement) || statements.length > 1;
+        }
+      }
+    }
+    return false;
+  }
+
+  @Override
+  protected void doUnwrap(final PsiElement element, final Context context) throws IncorrectOperationException {
+    final PyWithStatement withStatement = (PyWithStatement)element;
+    context.extractPart(withStatement);
+    context.delete(withStatement);
+  }
+}
+

python/testData/refactoring/unwrap/withUnwrap_after.py

@@ -0,0 +1 @@
+pr<caret>int(2)

python/testData/refactoring/unwrap/withUnwrap_before.py

@@ -0,0 +1,2 @@
+with locking(myLock):
+        pr<caret>int(2)

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

@@ -5,6 +5,7 @@ import com.intellij.openapi.actionSystem.AnAction;
 import com.intellij.openapi.editor.Editor;
 import com.intellij.psi.PsiFile;
 import com.jetbrains.python.fixtures.PyTestCase;
+import com.jetbrains.python.psi.LanguageLevel;
 
 import java.util.List;
 
@@ -37,10 +38,22 @@ public class PyUnwrapperTest extends PyTestCase {
 
   public void testForUnwrap()                         throws Throwable {doTest();}
 
+  public void testWithUnwrap()                        throws Throwable {doTest(LanguageLevel.PYTHON32);}
+
   private void doTest() {
     doTest(0);
   }
 
+  private void doTest(LanguageLevel languageLevel) {
+    setLanguageLevel(languageLevel);
+    try {
+      doTest(0);
+    }
+    finally {
+      setLanguageLevel(null);
+    }
+  }
+
   private void doTest(final int option) {
     String before = "refactoring/unwrap/" + getTestName(true) + "_before.py";
     String after = "refactoring/unwrap/" + getTestName(true) + "_after.py";