PY-1842 Extract method within while breaks content and generates endless loop

python/src/com/jetbrains/python/codeInsight/codeFragment/PyCodeFragmentBuilder.java

@@ -31,8 +31,23 @@ public class PyCodeFragmentBuilder extends PyRecursiveElementVisitor {
     endOffset = end;
   }
 
+  @Override
+  public void visitPyAssignmentStatement(final PyAssignmentStatement node) {
+    final PyExpression value = node.getAssignedValue();
+    if (value != null) {
+      value.accept(this);
+    }
+    for (PyExpression expression : node.getTargets()) {
+      expression.accept(this);
+    }
+  }
+
   @Override
   public void visitPyAugAssignmentStatement(final PyAugAssignmentStatement node) {
+    final PyExpression value = node.getValue();
+    if (value != null) {
+      value.accept(this);
+    }
     final PyExpression target = node.getTarget();
     if (target instanceof PyReferenceExpression){
       visitPyReferenceExpression((PyReferenceExpression) target);
@@ -40,6 +55,7 @@ public class PyCodeFragmentBuilder extends PyRecursiveElementVisitor {
     }
   }
 
+
   @Override
   public void visitPyTargetExpression(final PyTargetExpression node) {
     processDeclaration(node);
@@ -144,6 +160,10 @@ public class PyCodeFragmentBuilder extends PyRecursiveElementVisitor {
     final String name = element.getName();
     // Collect in variables
     if (position == Position.INSIDE) {
+      // support declarations within loops
+      if (inElements.contains(name) && PsiTreeUtil.getParentOfType(element, PyLoopStatement.class) != null){
+        outElements.add(name);
+      }
       // Add modification inside
       List<PyElement> list = modifiedInsideMap.get(name);
       if (list == null) {

python/testData/codeInsight/codefragment/while.test

@@ -0,0 +1,14 @@
+def fib(n):
+    a, b = 0, 1
+    while b < n:
+        <begin>print(b, end=' ')
+        a, b = b, a+b<end>
+    print()
+fib(7)
+<result>
+In:
+a
+b
+Out:
+a
+b

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

@@ -110,6 +110,10 @@ public class PyCodeFragmentTest extends LightMarkedTestCase {
     doTest();
   }
 
+  public void testWhile() throws Exception {
+    doTest();
+  }
+
   public void testEmpty() throws Exception {
     doTest();
   }