PY-5475 Don't attempt to inline expression if it contains comments

python/src/com/jetbrains/python/refactoring/introduce/IntroduceHandler.java

@@ -31,10 +31,7 @@ import com.intellij.openapi.project.Project;
 import com.intellij.openapi.util.Pair;
 import com.intellij.openapi.util.Pass;
 import com.intellij.openapi.util.TextRange;
-import com.intellij.psi.PsiElement;
-import com.intellij.psi.PsiFile;
-import com.intellij.psi.PsiWhiteSpace;
-import com.intellij.psi.TokenType;
+import com.intellij.psi.*;
 import com.intellij.psi.util.PsiTreeUtil;
 import com.intellij.refactoring.IntroduceTargetChooser;
 import com.intellij.refactoring.RefactoringActionHandler;
@@ -511,9 +508,7 @@ abstract public class IntroduceHandler implements RefactoringActionHandler {
   public PyAssignmentStatement createDeclaration(IntroduceOperation operation) {
     final Project project = operation.getProject();
     final PyExpression initializer = operation.getInitializer();
-    InitializerTextBuilder builder = new InitializerTextBuilder();
-    initializer.accept(builder);
-    String assignmentText = operation.getName() + " = " + builder.result();
+    String assignmentText = operation.getName() + " = " + new InitializerTextBuilder(initializer).result();
     PsiElement anchor = operation.isReplaceAll()
                         ? findAnchor(operation.getOccurrences())
                         : PsiTreeUtil.getParentOfType(initializer, PyStatement.class);
@@ -523,6 +518,18 @@ abstract public class IntroduceHandler implements RefactoringActionHandler {
   private static class InitializerTextBuilder extends PyRecursiveElementVisitor {
     private final StringBuilder myResult = new StringBuilder();
 
+    public InitializerTextBuilder(@NotNull PyExpression expression) {
+      if (PsiTreeUtil.findChildOfType(expression, PsiComment.class) != null) {
+        myResult.append(expression.getText());
+      }
+      else {
+        expression.accept(this);
+      }
+      if (needToWrapTopLevelExpressionInParenthesis(expression)) {
+        myResult.insert(0, "(").append(")");
+      }
+    }
+
     @Override
     public void visitWhiteSpace(PsiWhiteSpace space) {
       myResult.append(space.getText().replace('\n', ' ').replace("\\", ""));
@@ -560,17 +567,6 @@ abstract public class IntroduceHandler implements RefactoringActionHandler {
       }
     }
 
-    @Override
-    public void visitPyGeneratorExpression(PyGeneratorExpression node) {
-      final PsiElement firstChild = node.getFirstChild();
-      if (firstChild != null && firstChild.getNode().getElementType() != PyTokenTypes.LPAR) {
-        myResult.append("(").append(node.getText()).append(")");
-      }
-      else {
-        super.visitPyGeneratorExpression(node);
-      }
-    }
-
     @Override
     public void visitElement(PsiElement element) {
       if (element.getChildren().length == 0) {
@@ -581,6 +577,16 @@ abstract public class IntroduceHandler implements RefactoringActionHandler {
       }
     }
 
+    private boolean needToWrapTopLevelExpressionInParenthesis(@NotNull PyExpression node) {
+      if (node instanceof PyGeneratorExpression) {
+        final PsiElement firstChild = node.getFirstChild();
+        if (firstChild != null && firstChild.getNode().getElementType() != PyTokenTypes.LPAR) {
+          return true;
+        }
+      }
+      return false;
+    }
+
     public String result() {
       return myResult.toString();
     }

python/testData/refactoring/introduceVariable/functionCallWithCommentNotInlined.after.py

@@ -0,0 +1,8 @@
+import subprocess as sp
+
+a = sp.check_output(
+    args=['python', '-c', 'print("Spam")'],
+    # read errors too
+    stderr=sp.STDOUT
+)
+print(a)

python/testData/refactoring/introduceVariable/functionCallWithCommentNotInlined.py

@@ -0,0 +1,7 @@
+import subprocess as sp
+
+print(sp.c<caret>heck_output(
+    args=['python', '-c', 'print("Spam")'],
+    # read errors too
+    stderr=sp.STDOUT
+))

python/testData/refactoring/introduceVariable/generatorExpressionWithCommentNotInlined.after.py

@@ -0,0 +1,8 @@
+baz = [1, 2]
+a = (
+    el  # comment
+    if el >= 0
+    else -el
+    for el in baz
+)
+foo = bar(*a)

python/testData/refactoring/introduceVariable/generatorExpressionWithCommentNotInlined.py

@@ -0,0 +1,7 @@
+baz = [1, 2]
+foo = bar(*<selection>(
+    el  # comment
+    if el >= 0
+    else -el
+    for el in baz
+)</selection>)

python/testData/refactoring/introduceVariable/noRedundantParenthesisAroundGeneratorExpression.after.py

@@ -0,0 +1,2 @@
+a = list(i for in range(100) if x % 3 == 0)
+xs = a

python/testData/refactoring/introduceVariable/noRedundantParenthesisAroundGeneratorExpression.py

@@ -0,0 +1,2 @@
+xs = <selection>list(i for in range(100)
+                     if x % 3 == 0)</selection>

python/testSrc/com/jetbrains/python/refactoring/PyIntroduceVariableTest.java

@@ -224,10 +224,16 @@ public class PyIntroduceVariableTest extends PyIntroduceTestCase {
     doTest();
   }
 
+  // PY-11909
   public void testGeneratorParameter() {
     doTest();
   }
 
+  // PY-11909
+  public void testNoRedundantParenthesisAroundGeneratorExpression() {
+    doTest();
+  }
+
   // PY-10964
   public void testMultiReference() {
     myFixture.configureByFile(getTestName(true) + ".py");
@@ -245,7 +251,19 @@ public class PyIntroduceVariableTest extends PyIntroduceTestCase {
     }
   }
 
-  public void testSelectionBreaksBinaryOperator() {doTest();}
+  // PY-5475
+  public void testGeneratorExpressionWithCommentNotInlined() {
+    doTest();
+  }
+
+  // PY-5475
+  public void testFunctionCallWithCommentNotInlined() {
+    doTest();
+  }
+
+  public void testSelectionBreaksBinaryOperator() {
+    doTest();
+  }
 
   private void doTestCannotPerform() {
     boolean thrownExpectedException = false;