diff --git a/python/src/com/jetbrains/python/PyBundle.properties b/python/src/com/jetbrains/python/PyBundle.properties index 61a2b4948728..2cfbf828ba55 100644 --- a/python/src/com/jetbrains/python/PyBundle.properties +++ b/python/src/com/jetbrains/python/PyBundle.properties @@ -221,6 +221,7 @@ surround.with.try.except.template=try / except ########################################################################################################################## unwrap.if=Unwrap if... unwrap.while=Unwrap while... +unwrap.for=Unwrap for... unwrap.try=Unwrap try... unwrap.else=Unwrap else... unwrap.elif=Unwrap elif... diff --git a/python/src/com/jetbrains/python/refactoring/unwrap/PyForUnwrapper.java b/python/src/com/jetbrains/python/refactoring/unwrap/PyForUnwrapper.java new file mode 100644 index 000000000000..64d7c4006682 --- /dev/null +++ b/python/src/com/jetbrains/python/refactoring/unwrap/PyForUnwrapper.java @@ -0,0 +1,34 @@ +package com.jetbrains.python.refactoring.unwrap; + +import com.intellij.psi.PsiElement; +import com.intellij.util.IncorrectOperationException; +import com.jetbrains.python.PyBundle; +import com.jetbrains.python.psi.*; + +/** + * User : ktisha + */ +public class PyForUnwrapper extends PyUnwrapper { + public PyForUnwrapper() { + super(PyBundle.message("unwrap.for")); + } + + public boolean isApplicableTo(PsiElement e) { + if (e instanceof PyForStatement) { + final PyStatementList statementList = ((PyForStatement)e).getForPart().getStatementList(); + 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 PyForStatement forStatement = (PyForStatement)element; + context.extractPart(forStatement); + context.delete(forStatement); + } +} + diff --git a/python/src/com/jetbrains/python/refactoring/unwrap/PyUnwrapDescriptor.java b/python/src/com/jetbrains/python/refactoring/unwrap/PyUnwrapDescriptor.java index 56bf8e1dccee..9de763b9ca61 100644 --- a/python/src/com/jetbrains/python/refactoring/unwrap/PyUnwrapDescriptor.java +++ b/python/src/com/jetbrains/python/refactoring/unwrap/PyUnwrapDescriptor.java @@ -16,7 +16,8 @@ public class PyUnwrapDescriptor extends UnwrapDescriptorBase{ new PyElseUnwrapper(), new PyElIfUnwrapper(), new PyElIfRemover(), - new PyTryUnwrapper() + new PyTryUnwrapper(), + new PyForUnwrapper() }; } } diff --git a/python/src/com/jetbrains/python/refactoring/unwrap/PyUnwrapper.java b/python/src/com/jetbrains/python/refactoring/unwrap/PyUnwrapper.java index 4d36a34f16fd..930d4f9191d7 100644 --- a/python/src/com/jetbrains/python/refactoring/unwrap/PyUnwrapper.java +++ b/python/src/com/jetbrains/python/refactoring/unwrap/PyUnwrapper.java @@ -65,6 +65,10 @@ public abstract class PyUnwrapper extends AbstractUnwrapper final PyTryPart part = ((PyTryExceptStatement)from).getTryPart(); statementList = part.getStatementList(); } + else if (from instanceof PyForStatement) { + final PyForPart part = ((PyForStatement)from).getForPart(); + statementList = part.getStatementList(); + } if (statementList != null) extract(statementList.getFirstChild(), statementList.getLastChild(), from); } diff --git a/python/testData/refactoring/unwrap/forUnwrap_after.py b/python/testData/refactoring/unwrap/forUnwrap_after.py new file mode 100644 index 000000000000..4606edb69574 --- /dev/null +++ b/python/testData/refactoring/unwrap/forUnwrap_after.py @@ -0,0 +1 @@ +print 1 diff --git a/python/testData/refactoring/unwrap/forUnwrap_before.py b/python/testData/refactoring/unwrap/forUnwrap_before.py new file mode 100644 index 000000000000..fa36f9ce94c2 --- /dev/null +++ b/python/testData/refactoring/unwrap/forUnwrap_before.py @@ -0,0 +1,2 @@ +for item in range(1): + print 1 \ 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 32cb0087fb66..961968992e78 100644 --- a/python/testSrc/com/jetbrains/python/refactoring/PyUnwrapperTest.java +++ b/python/testSrc/com/jetbrains/python/refactoring/PyUnwrapperTest.java @@ -35,6 +35,8 @@ public class PyUnwrapperTest extends PyTestCase { public void testTryUnwrap() throws Throwable {doTest();} + public void testForUnwrap() throws Throwable {doTest();} + private void doTest() { doTest(0); }