From 22ec0dd9b342e9622f50984069d52386fba6c06b Mon Sep 17 00:00:00 2001
From: Ekaterina Tuzova <Ekaterina.Tuzova@jetbrains.com>
Date: Fri, 22 Feb 2013 15:06:13 +0400
Subject: [PATCH] Keep else part on unwrap for and while

fixed  PY-8569 Unwrap for loop statement should not delete else part
---
 .../jetbrains/python/refactoring/unwrap/PyForUnwrapper.java   | 1 +
 .../jetbrains/python/refactoring/unwrap/PyWhileUnwrapper.java | 1 +
 python/testData/refactoring/unwrap/forElseUnwrap_after.py     | 2 ++
 python/testData/refactoring/unwrap/forElseUnwrap_before.py    | 4 ++++
 python/testData/refactoring/unwrap/whileElseUnwrap_after.py   | 2 ++
 python/testData/refactoring/unwrap/whileElseUnwrap_before.py  | 4 ++++
 .../com/jetbrains/python/refactoring/PyUnwrapperTest.java     | 2 ++
 7 files changed, 16 insertions(+)
 create mode 100644 python/testData/refactoring/unwrap/forElseUnwrap_after.py
 create mode 100644 python/testData/refactoring/unwrap/forElseUnwrap_before.py
 create mode 100644 python/testData/refactoring/unwrap/whileElseUnwrap_after.py
 create mode 100644 python/testData/refactoring/unwrap/whileElseUnwrap_before.py

diff --git a/python/src/com/jetbrains/python/refactoring/unwrap/PyForUnwrapper.java b/python/src/com/jetbrains/python/refactoring/unwrap/PyForUnwrapper.java
index 64d7c4006682..099f71ad9c08 100644
--- a/python/src/com/jetbrains/python/refactoring/unwrap/PyForUnwrapper.java
+++ b/python/src/com/jetbrains/python/refactoring/unwrap/PyForUnwrapper.java
@@ -28,6 +28,7 @@ public class PyForUnwrapper extends PyUnwrapper {
   protected void doUnwrap(final PsiElement element, final Context context) throws IncorrectOperationException {
     final PyForStatement forStatement = (PyForStatement)element;
     context.extractPart(forStatement);
+    context.extractPart(forStatement.getElsePart());
     context.delete(forStatement);
   }
 }
diff --git a/python/src/com/jetbrains/python/refactoring/unwrap/PyWhileUnwrapper.java b/python/src/com/jetbrains/python/refactoring/unwrap/PyWhileUnwrapper.java
index b9fb2b72e6cf..55bbb8aad768 100644
--- a/python/src/com/jetbrains/python/refactoring/unwrap/PyWhileUnwrapper.java
+++ b/python/src/com/jetbrains/python/refactoring/unwrap/PyWhileUnwrapper.java
@@ -28,6 +28,7 @@ public class PyWhileUnwrapper extends PyUnwrapper {
   protected void doUnwrap(final PsiElement element, final Context context) throws IncorrectOperationException {
     final PyWhileStatement whileStatement = (PyWhileStatement)element;
     context.extractPart(whileStatement);
+    context.extractPart(whileStatement.getElsePart());
     context.delete(whileStatement);
   }
 }
diff --git a/python/testData/refactoring/unwrap/forElseUnwrap_after.py b/python/testData/refactoring/unwrap/forElseUnwrap_after.py
new file mode 100644
index 000000000000..474b42e9c37a
--- /dev/null
+++ b/python/testData/refactoring/unwrap/forElseUnwrap_after.py
@@ -0,0 +1,2 @@
+print 1
+print 2
diff --git a/python/testData/refactoring/unwrap/forElseUnwrap_before.py b/python/testData/refactoring/unwrap/forElseUnwrap_before.py
new file mode 100644
index 000000000000..572a1de74d30
--- /dev/null
+++ b/python/testData/refactoring/unwrap/forElseUnwrap_before.py
@@ -0,0 +1,4 @@
+for item in range(1):
+        prin<caret>t 1
+else:
+        print 2
\ No newline at end of file
diff --git a/python/testData/refactoring/unwrap/whileElseUnwrap_after.py b/python/testData/refactoring/unwrap/whileElseUnwrap_after.py
new file mode 100644
index 000000000000..7107e31922ff
--- /dev/null
+++ b/python/testData/refactoring/unwrap/whileElseUnwrap_after.py
@@ -0,0 +1,2 @@
+x = 1
+print 2
diff --git a/python/testData/refactoring/unwrap/whileElseUnwrap_before.py b/python/testData/refactoring/unwrap/whileElseUnwrap_before.py
new file mode 100644
index 000000000000..0fcb02a4bbe5
--- /dev/null
+++ b/python/testData/refactoring/unwrap/whileElseUnwrap_before.py
@@ -0,0 +1,4 @@
+while True:
+    x =<caret> 1
+else:
+    print 2
\ 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 5ea39abf3b33..7c7b72371fab 100644
--- a/python/testSrc/com/jetbrains/python/refactoring/PyUnwrapperTest.java
+++ b/python/testSrc/com/jetbrains/python/refactoring/PyUnwrapperTest.java
@@ -21,6 +21,7 @@ public class PyUnwrapperTest extends PyTestCase {
   public void testWhileUnwrap()                       throws Throwable {doTest();}
   public void testWhileUnwrapEmpty()                  throws Throwable {doNegativeTest();}
   public void testWhileUnwrapMultipleStatements()     throws Throwable {doTest();}
+  public void testWhileElseUnwrap()                   throws Throwable {doTest();}
 
   public void testIfWithElseUnwrap()                  throws Throwable {doTest();}
   public void testIfInWhileUnwrap()                   throws Throwable {doTest();}
@@ -40,6 +41,7 @@ public class PyUnwrapperTest extends PyTestCase {
   public void testTryElseFinallyUnwrap()              throws Throwable {doTest();}
 
   public void testForUnwrap()                         throws Throwable {doTest();}
+  public void testForElseUnwrap()                     throws Throwable {doTest();}
 
   public void testWithUnwrap()                        throws Throwable {doTest(LanguageLevel.PYTHON32);}