PY-25001 Implement Complete Current Statement for collection literals

It works by automatically inserting a trailing comma before a line break in
multiline collection literals. Additionally, a colon is inserted between a key
and a value in dict literals.

Because of the language ambiguity regarding syntax of some incomplete collection
literals, colon is not inserted after the first key of a dict literal
(it's indistinguishable from a set literal with one item), and comma is not
inserted after the first item of a parenthesized tuple (it's indistinguishable
from a parenthesized expression).

The idea was taken from such implementation of Complete Current Statement for
JSON.

GitOrigin-RevId: 49ff9857d8c12476bfa9aacaed0b49faa99810fd

python/python-psi-api/src/com/jetbrains/python/psi/impl/PyPsiUtils.java

@@ -98,7 +98,7 @@ public final class PyPsiUtils {
   }
 
   /**
-   * Finds first non-whitespace sibling after given PSI element but stops at first whitespace containing line feed.
+   * Returns the first non-whitespace sibling following the given element but within its line boundaries.
    */
   @Nullable
   public static PsiElement getNextNonWhitespaceSiblingOnSameLine(@NotNull PsiElement element) {
@@ -115,6 +115,24 @@ public final class PyPsiUtils {
     return null;
   }
 
+  /**
+   * Returns the first non-whitespace sibling preceding the given element but within its line boundaries.
+   */
+  @Nullable
+  public static PsiElement getPrevNonWhitespaceSiblingOnSameLine(@NotNull PsiElement element) {
+    PsiElement cur = element.getPrevSibling();
+    while (cur != null) {
+      if (!(cur instanceof PsiWhiteSpace)) {
+        return cur;
+      }
+      else if (cur.textContains('\n')) {
+        break;
+      }
+      cur = cur.getPrevSibling();
+    }
+    return null;
+  }
+
   /**
    * Finds first non-whitespace sibling after given AST node.
    */

python/src/com/jetbrains/python/codeInsight/editorActions/smartEnter/PySmartEnterProcessor.java

@@ -48,6 +48,7 @@ public class PySmartEnterProcessor extends SmartEnterProcessor {
     .add(new PyFunctionFixer())
     .add(new PyClassFixer())
     .add(new PyWithFixer())
+    .add(new PyCollectionLiteralFixer())
     .build();
   private static final List<EnterProcessor> ourProcessors = ImmutableList.of(new PyCommentBreakerEnterProcessor(),
                                                                              new PyPlainEnterProcessor());

python/src/com/jetbrains/python/codeInsight/editorActions/smartEnter/fixers/PyCollectionLiteralFixer.java

@@ -0,0 +1,88 @@
+// Copyright 2000-2020 JetBrains s.r.o. Use of this source code is governed by the Apache 2.0 license that can be found in the LICENSE file.
+package com.jetbrains.python.codeInsight.editorActions.smartEnter.fixers;
+
+import com.intellij.openapi.editor.Document;
+import com.intellij.openapi.editor.Editor;
+import com.intellij.psi.*;
+import com.intellij.util.ArrayUtil;
+import com.intellij.util.ObjectUtils;
+import com.jetbrains.python.PyPsiBundle;
+import com.jetbrains.python.codeInsight.editorActions.smartEnter.PySmartEnterProcessor;
+import com.jetbrains.python.codeInsight.editorActions.smartEnter.SmartEnterUtil;
+import com.jetbrains.python.psi.*;
+import com.jetbrains.python.psi.impl.PyPsiUtils;
+import org.jetbrains.annotations.NotNull;
+import org.jetbrains.annotations.Nullable;
+
+import static com.jetbrains.python.psi.PyUtil.as;
+
+public class PyCollectionLiteralFixer extends PyFixer<PySequenceExpression> {
+  public PyCollectionLiteralFixer() {
+    super(PySequenceExpression.class);
+  }
+
+  @Override
+  protected boolean isApplicable(@NotNull Editor editor, @NotNull PySequenceExpression element) {
+    return true;
+  }
+
+  @Override
+  protected void doApply(@NotNull Editor editor, @NotNull PySmartEnterProcessor processor, @NotNull PySequenceExpression collection) {
+    int caretOffset = editor.getCaretModel().getOffset();
+    PsiElement collectionItemAnchor = collection.getContainingFile().findElementAt(caretOffset);
+    assert collectionItemAnchor != null;
+    if (collectionItemAnchor instanceof PsiWhiteSpace) {
+      collectionItemAnchor = PyPsiUtils.getPrevNonWhitespaceSiblingOnSameLine(collectionItemAnchor);
+    }
+    if (collectionItemAnchor != null && isMissingColonError(collectionItemAnchor)) {
+      collectionItemAnchor = collectionItemAnchor.getPrevSibling();
+    }
+    // Since surrounding parentheses don't belong to PyTupleExpression expression,
+    // whitespace after its last element is a sibling of the tuple itself
+    if (collection instanceof PyTupleExpression && collectionItemAnchor == collection) {
+      collectionItemAnchor = ArrayUtil.getLastElement(collection.getElements());
+    }
+    if (collectionItemAnchor == null) {
+      return;
+    }
+    PsiElement collectionItem = PyPsiUtils.getParentRightBefore(collectionItemAnchor, collection);
+    if (!(collectionItem instanceof PyExpression)) {
+      return;
+    }
+
+    Document document = editor.getDocument();
+    int collectionItemLastLine = document.getLineNumber(collectionItem.getTextRange().getEndOffset());
+    int caretLine = document.getLineNumber(caretOffset);
+    if (collectionItemLastLine != caretLine) {
+      return;
+    }
+
+    PsiElement nextOnSameLine = PyPsiUtils.getNextNonWhitespaceSiblingOnSameLine(collectionItem);
+    if (nextOnSameLine != null && isMissingColonError(nextOnSameLine)) {
+      nextOnSameLine = PyPsiUtils.getNextNonWhitespaceSiblingOnSameLine(nextOnSameLine);
+    }
+    if (nextOnSameLine == null || nextOnSameLine instanceof PsiComment) {
+      int separatorOffset = collectionItem.getTextRange().getEndOffset();
+      PyKeyValueExpression keyValuePair = as(collectionItem, PyKeyValueExpression.class);
+      if (collection instanceof PyDictLiteralExpression && keyValuePair == null && !(collectionItem instanceof PyDoubleStarExpression)) {
+        document.insertString(separatorOffset, ": ");
+        processor.registerUnresolvedError(separatorOffset + 2);
+      }
+      // Our parser can't handle "'key': ," sequence
+      else if (!(keyValuePair != null && keyValuePair.getValue() == null)) {
+        document.insertString(separatorOffset, ",");
+        editor.getCaretModel().moveToOffset(separatorOffset + 1);
+        SmartEnterUtil.plainEnter(editor);
+        // The default behavior is to move the caret after the containing statement, unless other caretOffset is specified explicitly
+        processor.registerUnresolvedError(editor.getCaretModel().getOffset());
+        // Forcibly commit the document to prevent other fixers and enter processors from being run
+        PsiDocumentManager.getInstance(collection.getProject()).commitDocument(document);
+      }
+    }
+  }
+
+  private static boolean isMissingColonError(@NotNull PsiElement element) {
+    PsiErrorElement errorElement = as(element, PsiErrorElement.class);
+    return errorElement != null && PyPsiBundle.message("PARSE.expected.colon").equals(errorElement.getErrorDescription());
+  }
+}

python/testData/codeInsight/smartEnter/multilineDictAfterUnpacking.py

@@ -0,0 +1,4 @@
+def func(**kwargs):
+    d = {
+        **kwargs<caret>
+    }

python/testData/codeInsight/smartEnter/multilineDictAfterUnpacking_after.py

@@ -0,0 +1,5 @@
+def func(**kwargs):
+    d = {
+        **kwargs,
+        <caret>
+    }

python/testData/codeInsight/smartEnter/multilineDictLiteralAfterKeyWithColonButNoValue.py

@@ -0,0 +1,3 @@
+d = {
+    'key': <caret>
+}

python/testData/codeInsight/smartEnter/multilineDictLiteralAfterKeyWithColonButNoValue_after.py

@@ -0,0 +1,4 @@
+d = {
+    'key':
+        <caret>
+}

python/testData/codeInsight/smartEnter/multilineDictLiteralAfterKeyWithoutColon.py

@@ -0,0 +1,4 @@
+d = {
+  'key1': 1,
+  'key2'<caret>
+}

python/testData/codeInsight/smartEnter/multilineDictLiteralAfterKeyWithoutColon_after.py

@@ -0,0 +1,4 @@
+d = {
+  'key1': 1,
+  'key2': <caret>
+}

python/testData/codeInsight/smartEnter/multilineDictLiteralInsideKeyWithoutColon.py

@@ -0,0 +1,4 @@
+d = {
+    'key1': 1,
+    'key2<caret>'
+}

python/testData/codeInsight/smartEnter/multilineDictLiteralInsideKeyWithoutColon_after.py

@@ -0,0 +1,4 @@
+d = {
+    'key1': 1,
+    'key2': <caret>
+}

python/testData/codeInsight/smartEnter/multilineDictLiteralValueOnOtherLine.py

@@ -0,0 +1,4 @@
+d = {
+    'key': <caret>
+        42
+}

python/testData/codeInsight/smartEnter/multilineDictLiteralValueOnOtherLine_after.py

@@ -0,0 +1,5 @@
+d = {
+    'key':
+    <caret>
+        42
+}

python/testData/codeInsight/smartEnter/multilineDictLiteralValueOnSameLine.py

@@ -0,0 +1,3 @@
+d = {
+    'key': 42<caret>
+}

python/testData/codeInsight/smartEnter/multilineDictLiteralValueOnSameLine_after.py

@@ -0,0 +1,4 @@
+d = {
+    'key': 42,
+    <caret>
+}

python/testData/codeInsight/smartEnter/multilineListLiteralAfterItem.py

@@ -0,0 +1,3 @@
+xs = [
+    1<caret>
+]

python/testData/codeInsight/smartEnter/multilineListLiteralAfterItem_after.py

@@ -0,0 +1,4 @@
+xs = [
+    1,
+    <caret>
+]

python/testData/codeInsight/smartEnter/multilineListLiteralEmptyLine.py

@@ -0,0 +1,3 @@
+xs = [
+    <caret>
+]

python/testData/codeInsight/smartEnter/multilineListLiteralEmptyLine_after.py

@@ -0,0 +1,4 @@
+xs = [
+
+    <caret>
+]

python/testData/codeInsight/smartEnter/multilineListLiteralInsideCommentFollowingItem.py

@@ -0,0 +1,3 @@
+xs = [
+    42  # com<caret>ment
+]

python/testData/codeInsight/smartEnter/multilineListLiteralInsideCommentFollowingItem_after.py

@@ -0,0 +1,4 @@
+xs = [
+    42  # comment
+    # <caret>
+]

python/testData/codeInsight/smartEnter/multilineListLiteralInsideCommentedLine.py

@@ -0,0 +1,3 @@
+xs = [
+    # com<caret>ment
+]

python/testData/codeInsight/smartEnter/multilineListLiteralInsideCommentedLine_after.py

@@ -0,0 +1,4 @@
+xs = [
+    # comment
+    # <caret>
+]

python/testData/codeInsight/smartEnter/multilineListLiteralInsideItem.py

@@ -0,0 +1,3 @@
+xs = [
+    4<caret>2
+]

python/testData/codeInsight/smartEnter/multilineListLiteralInsideItem_after.py

@@ -0,0 +1,4 @@
+xs = [
+    42,
+    <caret>
+]

python/testData/codeInsight/smartEnter/multilineListLiteralInsideMultilineItem.py

@@ -0,0 +1,5 @@
+xs = [
+    oct(
+        42<caret>
+    )
+]

python/testData/codeInsight/smartEnter/multilineListLiteralInsideMultilineItem_after.py

@@ -0,0 +1,6 @@
+xs = [
+    oct(
+        42
+        <caret>
+    )
+]

python/testData/codeInsight/smartEnter/multilineListLiteralItemFollowedByComment.py

@@ -0,0 +1,3 @@
+xs = [
+    42<caret>  # comment
+]

python/testData/codeInsight/smartEnter/multilineListLiteralItemFollowedByComment_after.py

@@ -0,0 +1,4 @@
+xs = [
+    42,  # comment
+    <caret>
+]

python/testData/codeInsight/smartEnter/multilineListLiteralItemFollowsOpeningBracket.py

@@ -0,0 +1,2 @@
+xs = ['foo'<caret>
+      ]

python/testData/codeInsight/smartEnter/multilineListLiteralItemFollowsOpeningBracket_after.py

@@ -0,0 +1,3 @@
+xs = ['foo',
+      <caret>
+      ]

python/testData/codeInsight/smartEnter/multilineListLiteralOnLastLineOfMultilineItem.py

@@ -0,0 +1,5 @@
+xs = [
+    oct(
+        42
+    <caret>)
+]

python/testData/codeInsight/smartEnter/multilineListLiteralOnLastLineOfMultilineItem_after.py

@@ -0,0 +1,6 @@
+xs = [
+    oct(
+        42
+    ),
+    <caret>
+]

python/testData/codeInsight/smartEnter/multilineSetLiteral.py

@@ -0,0 +1,3 @@
+xs = {
+    1<caret>
+}

python/testData/codeInsight/smartEnter/multilineSetLiteral_after.py

@@ -0,0 +1,4 @@
+xs = {
+    1,
+    <caret>
+}

python/testData/codeInsight/smartEnter/multilineTupleLiteralIntermediateElement.py

@@ -0,0 +1,5 @@
+xs = (
+    1,
+    2<caret>
+    3
+)

python/testData/codeInsight/smartEnter/multilineTupleLiteralIntermediateElement_after.py

@@ -0,0 +1,6 @@
+xs = (
+    1,
+    2,
+    <caret>
+    3
+)

python/testData/codeInsight/smartEnter/multilineTupleLiteralLastElement.py

@@ -0,0 +1,4 @@
+xs = (
+    1,
+    2 <caret>
+)

python/testData/codeInsight/smartEnter/multilineTupleLiteralLastElement_after.py

@@ -0,0 +1,5 @@
+xs = (
+    1,
+    2,
+    <caret>
+)

python/testSrc/com/jetbrains/python/PySmartEnterTest.java

@@ -117,6 +117,96 @@ public class PySmartEnterTest extends PyTestCase {
     doTest();
   }
 
+  // PY-25001
+  public void testMultilineListLiteralInsideItem() {
+    doTest();
+  }
+
+  // PY-25001
+  public void testMultilineListLiteralAfterItem() {
+    doTest();
+  }
+
+  // PY-25001
+  public void testMultilineListLiteralInsideMultilineItem() {
+    doTest();
+  }
+
+  // PY-25001
+  public void testMultilineListLiteralOnLastLineOfMultilineItem() {
+    doTest();
+  }
+
+  // PY-25001
+  public void testMultilineListLiteralInsideCommentFollowingItem() {
+    doTest();
+  }
+
+  // PY-25001
+  public void testMultilineListLiteralEmptyLine() {
+    doTest();
+  }
+
+  // PY-25001
+  public void testMultilineListLiteralInsideCommentedLine() {
+    doTest();
+  }
+
+  // PY-25001
+  public void testMultilineListLiteralItemFollowedByComment() {
+    doTest();
+  }
+
+  // PY-25001
+  public void testMultilineListLiteralItemFollowsOpeningBracket() {
+    doTest();
+  }
+
+  // PY-25001
+  public void testMultilineSetLiteral() {
+    doTest();
+  }
+
+  // PY-25001
+  public void testMultilineTupleLiteralLastElement() {
+    doTest();
+  }
+
+  // PY-25001
+  public void testMultilineTupleLiteralIntermediateElement() {
+    doTest();
+  }
+
+  // PY-25001
+  public void testMultilineDictLiteralAfterKeyWithoutColon() {
+    doTest();
+  }
+
+  // PY-25001
+  public void testMultilineDictLiteralInsideKeyWithoutColon() {
+    doTest();
+  }
+
+  // PY-25001
+  public void testMultilineDictLiteralAfterKeyWithColonButNoValue() {
+    doTest();
+  }
+
+  // PY-25001
+  public void testMultilineDictLiteralValueOnSameLine() {
+    doTest();
+  }
+
+  // PY-25001
+  public void testMultilineDictLiteralValueOnOtherLine() {
+    doTest();
+  }
+
+  // PY-25001
+  public void testMultilineDictAfterUnpacking() {
+    doTest();
+  }
+
   public void testDocRest() {
     CodeInsightSettings codeInsightSettings = CodeInsightSettings.getInstance();
     boolean oldStubOnEnter = codeInsightSettings.JAVADOC_STUB_ON_ENTER;