diff --git a/python/src/com/jetbrains/python/inspections/PyStatementEffectInspection.java b/python/src/com/jetbrains/python/inspections/PyStatementEffectInspection.java
index 0cc7a4abd140..a521ed76d92f 100644
--- a/python/src/com/jetbrains/python/inspections/PyStatementEffectInspection.java
+++ b/python/src/com/jetbrains/python/inspections/PyStatementEffectInspection.java
@@ -72,9 +72,6 @@ public class PyStatementEffectInspection extends PyInspection {
final PyTryPart tryPart = PsiTreeUtil.getParentOfType(node, PyTryPart.class);
if (tryPart != null) {
final PyStatementList statementList = tryPart.getStatementList();
- if (statementList == null) {
- return;
- }
if (statementList.getStatements().length == 1 && statementList.getStatements()[0] == node) {
return;
}
diff --git a/python/src/com/jetbrains/python/inspections/quickfix/StatementEffectFunctionCallQuickFix.java b/python/src/com/jetbrains/python/inspections/quickfix/StatementEffectFunctionCallQuickFix.java
index 95daf466172e..d852fbc55c83 100644
--- a/python/src/com/jetbrains/python/inspections/quickfix/StatementEffectFunctionCallQuickFix.java
+++ b/python/src/com/jetbrains/python/inspections/quickfix/StatementEffectFunctionCallQuickFix.java
@@ -18,14 +18,12 @@ 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.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.psi.*;
import org.jetbrains.annotations.NotNull;
/**
@@ -77,14 +75,23 @@ public class StatementEffectFunctionCallQuickFix implements LocalQuickFix {
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=' '");
next.delete();
}
stringBuilder.append(")");
+ if (commentText != null) {
+ stringBuilder.append(commentText);
+ }
expression.replace(elementGenerator.createFromText(LanguageLevel.forElement(expression), PyExpression.class,
stringBuilder.toString()));
}
diff --git a/python/testData/inspections/ReplacePrintComment.py b/python/testData/inspections/ReplacePrintComment.py
new file mode 100644
index 000000000000..1cf3778b69b1
--- /dev/null
+++ b/python/testData/inspections/ReplacePrintComment.py
@@ -0,0 +1 @@
+print '%s %s %s %s' % bar # <- doesn't work either
\ No newline at end of file
diff --git a/python/testData/inspections/ReplacePrintComment_after.py b/python/testData/inspections/ReplacePrintComment_after.py
new file mode 100644
index 000000000000..416deb8fd16e
--- /dev/null
+++ b/python/testData/inspections/ReplacePrintComment_after.py
@@ -0,0 +1 @@
+print('%s %s %s %s' % bar) # <- doesn't work either
diff --git a/python/testData/inspections/ReplacePrintEnd.py b/python/testData/inspections/ReplacePrintEnd.py
new file mode 100644
index 000000000000..afc37fbe2840
--- /dev/null
+++ b/python/testData/inspections/ReplacePrintEnd.py
@@ -0,0 +1 @@
+print var,
\ No newline at end of file
diff --git a/python/testData/inspections/ReplacePrintEnd_after.py b/python/testData/inspections/ReplacePrintEnd_after.py
new file mode 100644
index 000000000000..bae6b0c068a7
--- /dev/null
+++ b/python/testData/inspections/ReplacePrintEnd_after.py
@@ -0,0 +1 @@
+print(var, end=' ')
diff --git a/python/testSrc/com/jetbrains/python/PyQuickFixTest.java b/python/testSrc/com/jetbrains/python/PyQuickFixTest.java
index c36372d9ba64..68387a2c9c2f 100644
--- a/python/testSrc/com/jetbrains/python/PyQuickFixTest.java
+++ b/python/testSrc/com/jetbrains/python/PyQuickFixTest.java
@@ -251,6 +251,23 @@ public class PyQuickFixTest extends PyTestCase {
PyBundle.message("QFIX.statement.effect.introduce.variable"), true, true);
}
+ public void testReplacePrintEnd() {
+ runWithLanguageLevel(LanguageLevel.PYTHON34, new Runnable() {
+ @Override
+ public void run() {
+ doInspectionTest("ReplacePrintEnd.py", PyStatementEffectInspection.class, PyBundle.message("QFIX.statement.effect"), true, true);
+ }});
+ }
+
+ public void testReplacePrintComment() {
+ runWithLanguageLevel(LanguageLevel.PYTHON34, new Runnable() {
+ @Override
+ public void run() {
+ doInspectionTest("ReplacePrintComment.py", PyStatementEffectInspection.class, PyBundle.message("QFIX.statement.effect"), true,
+ true);
+ }});
+ }
+
public void testUnresolvedWith() { // PY-2083
setLanguageLevel(LanguageLevel.PYTHON25);
doInspectionTest("UnresolvedWith.py", PyUnresolvedReferencesInspection.class,