diff --git a/python/src/com/jetbrains/python/codeInsight/codeFragment/PyCodeFragmentBuilder.java b/python/src/com/jetbrains/python/codeInsight/codeFragment/PyCodeFragmentBuilder.java index de1ab103c5ac..6656aabcca42 100644 --- a/python/src/com/jetbrains/python/codeInsight/codeFragment/PyCodeFragmentBuilder.java +++ b/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 list = modifiedInsideMap.get(name); if (list == null) { diff --git a/python/testData/codeInsight/codefragment/while.test b/python/testData/codeInsight/codefragment/while.test new file mode 100644 index 000000000000..6304064dc37b --- /dev/null +++ b/python/testData/codeInsight/codefragment/while.test @@ -0,0 +1,14 @@ +def fib(n): + a, b = 0, 1 + while b < n: + print(b, end=' ') + a, b = b, a+b + print() +fib(7) + +In: +a +b +Out: +a +b diff --git a/python/testSrc/com/jetbrains/python/refactoring/PyCodeFragmentTest.java b/python/testSrc/com/jetbrains/python/refactoring/PyCodeFragmentTest.java index db7a65fa4a33..debeec5c747a 100644 --- a/python/testSrc/com/jetbrains/python/refactoring/PyCodeFragmentTest.java +++ 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(); }