diff --git a/python/src/com/jetbrains/python/codeInsight/intentions/ReplaceListComprehensionWithForIntention.java b/python/src/com/jetbrains/python/codeInsight/intentions/ReplaceListComprehensionWithForIntention.java index b66325e7fd77..dd2f13ba556c 100644 --- a/python/src/com/jetbrains/python/codeInsight/intentions/ReplaceListComprehensionWithForIntention.java +++ b/python/src/com/jetbrains/python/codeInsight/intentions/ReplaceListComprehensionWithForIntention.java @@ -127,15 +127,27 @@ public class ReplaceListComprehensionWithForIntention implements IntentionAction ComprhForComponent comp = forComps.get(0); PyForStatement pyForStatement = elementGenerator.createFromText(LanguageLevel.forElement(expression), PyForStatement.class, "for " + comp.getIteratorVariable().getText() + " in "+ comp.getIteratedList().getText() + ":\n a+1"); + List ifComps = ((PyListCompExpression)expression).getIfComponents(); if (ifComps.size() != 0) { addIfComponents(pyForStatement, ifComps, elementGenerator); } + PyStatement toReplace = pyForStatement.getForPart().getStatementList().getStatements()[0]; + while (toReplace instanceof PyIfStatement) { + toReplace = ((PyIfStatement)toReplace).getIfPart().getStatementList().getStatements()[0]; + } + toReplace.replace(statement); addForComponents(pyForStatement, ((PyListCompExpression)expression).getResultExpression(), elementGenerator, result); - pyStatement.replace(pyForStatement); + statement.getForPart().replace(pyForStatement); } } else { + while (pyStatement instanceof PyForStatement) { + pyStatement = ((PyForStatement)pyStatement).getForPart().getStatementList().getStatements()[0]; + while (pyStatement instanceof PyIfStatement) { + pyStatement = ((PyIfStatement)pyStatement).getIfPart().getStatementList().getStatements()[0]; + } + } pyStatement.replace(result); } } diff --git a/python/testData/intentions/afterReplaceListComprehensionWithFor.py b/python/testData/intentions/afterReplaceListComprehensionWithFor.py index 5bf49d51dbdf..017021e9affe 100644 --- a/python/testData/intentions/afterReplaceListComprehensionWithFor.py +++ b/python/testData/intentions/afterReplaceListComprehensionWithFor.py @@ -1,14 +1,14 @@ a = 2 smth = [] -for z in xrange(14): - if z != 3: - if z != 8: +for x in xrange(10): + if x != 2: + if x != 5: for y in xrange(12): if y != 1: if y != 4: - for x in xrange(10): - if x != 2: - if x != 5: + for z in xrange(14): + if z != 3: + if z != 8: smth.append(x + y + z) a = 5 \ No newline at end of file