diff --git a/python/src/com/jetbrains/python/inspections/quickfix/ConvertDictCompQuickFix.java b/python/src/com/jetbrains/python/inspections/quickfix/ConvertDictCompQuickFix.java index 45f20bed24a2..67a2bddfc079 100644 --- a/python/src/com/jetbrains/python/inspections/quickfix/ConvertDictCompQuickFix.java +++ b/python/src/com/jetbrains/python/inspections/quickfix/ConvertDictCompQuickFix.java @@ -52,16 +52,38 @@ public class ConvertDictCompQuickFix implements LocalQuickFix { } private static void replaceComprehension(Project project, PyDictCompExpression expression) { - List forComponents = expression.getForComponents(); if (expression.getResultExpression() instanceof PyKeyValueExpression) { - PyKeyValueExpression keyValueExpression = (PyKeyValueExpression)expression.getResultExpression(); - PyElementGenerator elementGenerator = PyElementGenerator.getInstance(project); + final PyKeyValueExpression keyValueExpression = (PyKeyValueExpression)expression.getResultExpression(); + final PyElementGenerator elementGenerator = PyElementGenerator.getInstance(project); assert keyValueExpression.getValue() != null; - expression.replace(elementGenerator.createFromText(LanguageLevel.getDefault(), PyExpressionStatement.class, - "dict([(" + keyValueExpression.getKey().getText() + ", " + - keyValueExpression.getValue().getText() + ") for " + - forComponents.get(0).getIteratorVariable().getText() + " in " + - forComponents.get(0).getIteratedList().getText() + "])")); + + final List components = expression.getComponents(); + final StringBuilder replacement = new StringBuilder("dict([(" + keyValueExpression.getKey().getText() + ", " + + keyValueExpression.getValue().getText() + ")"); + int slashNum = 1; + for (ComprehensionComponent component : components) { + if (component instanceof ComprhForComponent) { + replacement.append("for "); + replacement.append(((ComprhForComponent)component).getIteratorVariable().getText()); + replacement.append(" in "); + replacement.append(((ComprhForComponent)component).getIteratedList().getText()); + replacement.append(" "); + } + if (component instanceof ComprhIfComponent) { + final PyExpression test = ((ComprhIfComponent)component).getTest(); + if (test != null) { + replacement.append("if "); + replacement.append(test.getText()); + replacement.append(" "); + } + } + for (int i = 0; i != slashNum; ++i) + replacement.append("\t"); + ++slashNum; + } + replacement.append("])"); + + expression.replace(elementGenerator.createFromText(LanguageLevel.getDefault(), PyExpressionStatement.class, replacement.toString())); } } diff --git a/python/testData/inspections/DictComprehensionToCall.py b/python/testData/inspections/DictComprehensionToCall.py new file mode 100644 index 000000000000..b360be79e87d --- /dev/null +++ b/python/testData/inspections/DictComprehensionToCall.py @@ -0,0 +1 @@ +var = {k: v for k, v in zip('abc', range(3)) if k % 2} \ No newline at end of file diff --git a/python/testData/inspections/DictComprehensionToCall_after.py b/python/testData/inspections/DictComprehensionToCall_after.py new file mode 100644 index 000000000000..a71f9ffd1e89 --- /dev/null +++ b/python/testData/inspections/DictComprehensionToCall_after.py @@ -0,0 +1 @@ +var = dict([(k, v) for k, v in zip('abc', range(3)) if k % 2]) \ No newline at end of file diff --git a/python/testSrc/com/jetbrains/python/PyQuickFixTest.java b/python/testSrc/com/jetbrains/python/PyQuickFixTest.java index 4083775e8a41..ad74c78338f2 100644 --- a/python/testSrc/com/jetbrains/python/PyQuickFixTest.java +++ b/python/testSrc/com/jetbrains/python/PyQuickFixTest.java @@ -342,6 +342,11 @@ public class PyQuickFixTest extends PyTestCase { PyBundle.message("QFIX.replace.function.set.with.literal"), true, true); } + public void testDictComprehensionToCall() { + doInspectionTest("DictComprehensionToCall.py", PyCompatibilityInspection.class, + PyBundle.message("INTN.convert.dict.comp.to"), true, true); + } + public void testDocstringParams() { //PY-3394 PyDocumentationSettings documentationSettings = PyDocumentationSettings.getInstance(myFixture.getModule()); documentationSettings.setFormat(DocStringFormat.EPYTEXT);