diff options
Diffstat (limited to 'python/src/com/jetbrains/python/inspections/quickfix/StatementEffectFunctionCallQuickFix.java')
-rw-r--r-- | python/src/com/jetbrains/python/inspections/quickfix/StatementEffectFunctionCallQuickFix.java | 96 |
1 files changed, 86 insertions, 10 deletions
diff --git a/python/src/com/jetbrains/python/inspections/quickfix/StatementEffectFunctionCallQuickFix.java b/python/src/com/jetbrains/python/inspections/quickfix/StatementEffectFunctionCallQuickFix.java index 95daf466172e..43f4f0394015 100644 --- a/python/src/com/jetbrains/python/inspections/quickfix/StatementEffectFunctionCallQuickFix.java +++ b/python/src/com/jetbrains/python/inspections/quickfix/StatementEffectFunctionCallQuickFix.java @@ -18,14 +18,14 @@ package com.jetbrains.python.inspections.quickfix; import com.intellij.codeInspection.LocalQuickFix; import com.intellij.codeInspection.ProblemDescriptor; import com.intellij.openapi.project.Project; +import com.intellij.psi.PsiComment; import com.intellij.psi.PsiElement; import com.intellij.psi.PsiWhiteSpace; +import com.intellij.psi.impl.source.tree.LeafPsiElement; import com.jetbrains.python.PyBundle; import com.jetbrains.python.PyNames; -import com.jetbrains.python.psi.LanguageLevel; -import com.jetbrains.python.psi.PyElementGenerator; -import com.jetbrains.python.psi.PyExpression; -import com.jetbrains.python.psi.PyReferenceExpression; +import com.jetbrains.python.PyTokenTypes; +import com.jetbrains.python.psi.*; import org.jetbrains.annotations.NotNull; /** @@ -48,15 +48,17 @@ public class StatementEffectFunctionCallQuickFix implements LocalQuickFix { PsiElement expression = descriptor.getPsiElement(); if (expression != null && expression.isWritable() && expression instanceof PyReferenceExpression) { final String expressionText = expression.getText(); - if (PyNames.PRINT.equals(expressionText) || PyNames.EXEC.equals(expressionText)) - replacePrintExec(expression); + if (PyNames.PRINT.equals(expressionText)) + replacePrint(expression); + else if (PyNames.EXEC.equals(expressionText)) + replaceExec(expression); else expression.replace(PyElementGenerator.getInstance(project).createCallExpression(LanguageLevel.forElement(expression), expressionText)); } } - private static void replacePrintExec(@NotNull final PsiElement expression) { + private static void replaceExec(@NotNull final PsiElement expression) { final PyElementGenerator elementGenerator = PyElementGenerator.getInstance(expression.getProject()); final String expressionText = expression.getText(); final StringBuilder stringBuilder = new StringBuilder(expressionText + " ("); @@ -77,14 +79,88 @@ public class StatementEffectFunctionCallQuickFix implements LocalQuickFix { RemoveUnnecessaryBackslashQuickFix.removeBackSlash(next); if (whiteSpace != null) whiteSpace.delete(); + if (next == null) { + stringBuilder.append(")"); + expression.replace(elementGenerator.createFromText(LanguageLevel.forElement(expression), PyExpression.class, + stringBuilder.toString())); + return; + } + if (next instanceof PyExpressionStatement) { + final PyExpression expr = ((PyExpressionStatement)next).getExpression(); + if (expr instanceof PyBinaryExpression) { + addInArguments(stringBuilder, (PyBinaryExpression)expr); + } + else if (expr instanceof PyTupleExpression) { + final PyExpression[] elements = ((PyTupleExpression)expr).getElements(); + if (elements.length > 1) { + if (elements[0] instanceof PyBinaryExpression) { + addInArguments(stringBuilder, (PyBinaryExpression)elements[0]); + } + stringBuilder.append(", "); + stringBuilder.append(elements[1].getText()); + } + } + else { + stringBuilder.append(next.getText()); + } + } + else { + stringBuilder.append(next.getText()); + } + next.delete(); + stringBuilder.append(")"); + expression.replace(elementGenerator.createFromText(LanguageLevel.forElement(expression), PyExpression.class, + stringBuilder.toString())); + } + + private static void addInArguments(@NotNull final StringBuilder stringBuilder, @NotNull final PyBinaryExpression binaryExpression) { + final PsiElement operator = binaryExpression.getPsiOperator(); + if (operator instanceof LeafPsiElement && ((LeafPsiElement)operator).getElementType() == PyTokenTypes.IN_KEYWORD) { + stringBuilder.append(binaryExpression.getLeftExpression().getText()); + stringBuilder.append(", "); + final PyExpression rightExpression = binaryExpression.getRightExpression(); + if (rightExpression != null) + stringBuilder.append(rightExpression.getText()); + } + } + + private static void replacePrint(@NotNull final PsiElement expression) { + final PyElementGenerator elementGenerator = PyElementGenerator.getInstance(expression.getProject()); + final String expressionText = expression.getText(); + final StringBuilder stringBuilder = new StringBuilder(expressionText + " ("); + + final PsiElement whiteSpace = expression.getContainingFile().findElementAt(expression.getTextOffset() + expression.getTextLength()); + PsiElement next = null; + if (whiteSpace instanceof PsiWhiteSpace) { + final String whiteSpaceText = whiteSpace.getText(); + if (!whiteSpaceText.contains("\n")) { + next = whiteSpace.getNextSibling(); + while (next instanceof PsiWhiteSpace && whiteSpaceText.contains("\\")) { + next = next.getNextSibling(); + } + } + } + else + next = whiteSpace; + + RemoveUnnecessaryBackslashQuickFix.removeBackSlash(next); + if (whiteSpace != null) whiteSpace.delete(); + String commentText = null; if (next != null) { - final String text = next.getText(); + final PsiElement lastChild = next.getLastChild(); + if (lastChild instanceof PsiComment) { + commentText = lastChild.getText(); + } + final String text = next instanceof PyExpressionStatement ? ((PyExpressionStatement)next).getExpression().getText() : next.getText(); + stringBuilder.append(text); - if (text.endsWith(",") && PyNames.PRINT.equals(expressionText)) - stringBuilder.append(" end=' '"); + if (text.endsWith(",")) stringBuilder.append(" end=' '"); next.delete(); } stringBuilder.append(")"); + if (commentText != null) { + stringBuilder.append(commentText); + } expression.replace(elementGenerator.createFromText(LanguageLevel.forElement(expression), PyExpression.class, stringBuilder.toString())); } |