PY-16553 Simpler solution - add missing comma through PSI

diff --git a/python/src/com/jetbrains/python/codeInsight/intentions/PyBaseConvertCollectionLiteralIntention.java b/python/src/com/jetbrains/python/codeInsight/intentions/PyBaseConvertCollectionLiteralIntention.java
index 44557bcd0e915a069c95fce2bf95043dc9f35539..75e2d79408eee09402383298917e6da87a8b9017 100644 (file)
--- a/python/src/com/jetbrains/python/codeInsight/intentions/PyBaseConvertCollectionLiteralIntention.java
+++ b/python/src/com/jetbrains/python/codeInsight/intentions/PyBaseConvertCollectionLiteralIntention.java
@@ -90,27 +90,39 @@ public abstract class PyBaseConvertCollectionLiteralIntention extends BaseIntent
     final PySequenceExpression literal = findCollectionLiteralUnderCaret(editor, file);
     assert literal != null;
 
-    final PsiElement replacedElement;
-    if (literal instanceof PyTupleExpression && literal.getParent() instanceof PyParenthesizedExpression) {
-      replacedElement = literal.getParent();
-    }
-    else {
-      replacedElement = literal;
-    }
+    final PsiElement replacedElement = wrapCollection(literal);
+    final PsiElement copy = prepareOriginalElementCopy(replacedElement.copy());
 
-    final TextRange contentRange = getRangeOfContentWithoutBraces(replacedElement);
+    final TextRange contentRange = getRangeOfContentWithoutBraces(copy);
+    final String contentToWrap = contentRange.substring(copy.getText());
     final PyElementGenerator elementGenerator = PyElementGenerator.getInstance(project);
-    final String contentToWrap = prepareContent(replacedElement, literal, contentRange);
     final PyExpression newLiteral = elementGenerator.createExpressionFromText(LanguageLevel.forElement(file),
                                                                               myLeftBrace + contentToWrap + myRightBrace);
     replacedElement.replace(newLiteral);
   }
 
   @NotNull
-  protected String prepareContent(@NotNull PsiElement replacedElement, 
-                                  @NotNull PySequenceExpression collection, 
-                                  @NotNull TextRange contentRange) {
-    return contentRange.substring(replacedElement.getText());
+  protected PsiElement prepareOriginalElementCopy(@NotNull PsiElement copy) {
+    return copy;
+  }
+
+  @NotNull
+  protected static PySequenceExpression unwrapCollection(@NotNull PsiElement literal) {
+    final PyParenthesizedExpression parenthesizedExpression = as(literal, PyParenthesizedExpression.class);
+    if (parenthesizedExpression != null) {
+      final PyExpression containedExpression = parenthesizedExpression.getContainedExpression();
+      assert containedExpression != null;
+      return (PyTupleExpression)containedExpression;
+    }
+    return (PySequenceExpression)literal;
+  }
+
+  @NotNull
+  protected static PsiElement wrapCollection(@NotNull PySequenceExpression literal) {
+    if (literal instanceof PyTupleExpression && literal.getParent() instanceof PyParenthesizedExpression) {
+      return literal.getParent();
+    }
+    return literal;
   }
 
   @NotNull

diff --git a/python/src/com/jetbrains/python/codeInsight/intentions/PyConvertLiteralToTupleIntention.java b/python/src/com/jetbrains/python/codeInsight/intentions/PyConvertLiteralToTupleIntention.java
index 56d68c7ac7752f8d05f89313bb05f66afebb8ec7..5a6c7ea55b7a245ce75861bf5dcd6275829fa73f 100644 (file)
--- a/python/src/com/jetbrains/python/codeInsight/intentions/PyConvertLiteralToTupleIntention.java
+++ b/python/src/com/jetbrains/python/codeInsight/intentions/PyConvertLiteralToTupleIntention.java
@@ -15,10 +15,10 @@
  */
 package com.jetbrains.python.codeInsight.intentions;
 
-import com.intellij.openapi.util.TextRange;
+import com.intellij.lang.ASTNode;
 import com.intellij.psi.PsiElement;
-import com.intellij.psi.tree.IElementType;
 import com.jetbrains.python.PyTokenTypes;
+import com.jetbrains.python.psi.PyElementGenerator;
 import com.jetbrains.python.psi.PyExpression;
 import com.jetbrains.python.psi.PySequenceExpression;
 import com.jetbrains.python.psi.PyTupleExpression;
@@ -33,43 +33,20 @@ public class PyConvertLiteralToTupleIntention extends PyBaseConvertCollectionLit
     super(PyTupleExpression.class, "tuple", "(", ")");
   }
 
-
   @NotNull
   @Override
-  protected String prepareContent(@NotNull PsiElement replacedElement, 
-                                  @NotNull PySequenceExpression collection, 
-                                  @NotNull TextRange contentRange) {
-    assert !(collection instanceof PyTupleExpression);
-
-    final String contentWithoutBraces = super.prepareContent(replacedElement, collection, contentRange);
-    
-    final PyExpression[] elements = collection.getElements();
-    if (elements.length != 1) {
-      return contentWithoutBraces;
-    }
-    
-    final PsiElement lastChild = collection.getLastChild();
-    boolean endsWithComma = false;
-    final IElementType lastChildType = lastChild.getNode().getElementType();
-    if (lastChildType == PyTokenTypes.COMMA) {
-      endsWithComma = true;
-    }
-    else if (PyTokenTypes.CLOSE_BRACES.contains(lastChildType)) {
-      final PsiElement prev = PyPsiUtils.getPrevNonWhitespaceSibling(lastChild);
-      if (prev != null && prev.getNode().getElementType() == PyTokenTypes.COMMA) {
-        endsWithComma = true;
+  protected PsiElement prepareOriginalElementCopy(@NotNull PsiElement copy) {
+    final PySequenceExpression sequenceExpression = unwrapCollection(copy);
+    final PyExpression[] elements = sequenceExpression.getElements();
+    if (elements.length == 1) {
+      final PyExpression onlyElement = elements[0];
+      final PsiElement next = PyPsiUtils.getNextNonCommentSibling(onlyElement, true);
+      if (next != null && next.getNode().getElementType() != PyTokenTypes.COMMA) {
+        final PyElementGenerator generator = PyElementGenerator.getInstance(copy.getProject());
+        final ASTNode anchor = onlyElement.getNode().getTreeNext();
+        sequenceExpression.getNode().addChild(generator.createComma(), anchor);
       }
     }
-    if (endsWithComma) {
-      return contentWithoutBraces;
-    }
-
-    final PyExpression singleElem = elements[0];
-    final int commaOffset = singleElem.getTextRange().getEndOffset() - replacedElement.getTextRange().getStartOffset();
-
-    final String wholeText = replacedElement.getText();
-    return wholeText.substring(contentRange.getStartOffset(), commaOffset) + 
-           "," + 
-           wholeText.substring(commaOffset, contentRange.getEndOffset());
+    return copy;
   }
 }

diff --git a/python/testData/intentions/PyConvertCollectionLiteralIntentionTest/convertOneElementListWithCommaAfterCommentToTuple.py b/python/testData/intentions/PyConvertCollectionLiteralIntentionTest/convertOneElementListWithCommaAfterCommentToTuple.py
new file mode 100644 (file)
index 0000000..1dfb656
--- /dev/null
+++ b/python/testData/intentions/PyConvertCollectionLiteralIntentionTest/convertOneElementListWithCommaAfterCommentToTuple.py
@@ -0,0 +1,4 @@
+[
+    42  # foo
+    ,
+]

diff --git a/python/testData/intentions/PyConvertCollectionLiteralIntentionTest/convertOneElementListWithCommaAfterCommentToTuple_after.py b/python/testData/intentions/PyConvertCollectionLiteralIntentionTest/convertOneElementListWithCommaAfterCommentToTuple_after.py
new file mode 100644 (file)
index 0000000..7eedda7
--- /dev/null
+++ b/python/testData/intentions/PyConvertCollectionLiteralIntentionTest/convertOneElementListWithCommaAfterCommentToTuple_after.py
@@ -0,0 +1,4 @@
+(
+    42  # foo
+    ,
+)

diff --git a/python/testSrc/com/jetbrains/python/intentions/PyConvertCollectionLiteralIntentionTest.java b/python/testSrc/com/jetbrains/python/intentions/PyConvertCollectionLiteralIntentionTest.java
index efc2e267cdad657fe335dbc7309f5c2b9bf534dd..daa2206d2e4dd981357ccdc491363a419abf6d59 100644 (file)
--- a/python/testSrc/com/jetbrains/python/intentions/PyConvertCollectionLiteralIntentionTest.java
+++ b/python/testSrc/com/jetbrains/python/intentions/PyConvertCollectionLiteralIntentionTest.java
@@ -123,4 +123,9 @@ public class PyConvertCollectionLiteralIntentionTest extends PyIntentionTestCase
   public void testConvertOneElementListWithCommentToTuple() {
     doIntentionTest(CONVERT_LIST_TO_TUPLE);
   }
+  
+  // PY-16553
+  public void testConvertOneElementListWithCommaAfterCommentToTuple() {
+    doIntentionTest(CONVERT_LIST_TO_TUPLE);
+  }
 }