PY-10972 Properly align elements of expressions without closing bracket on enter

instead of using artificially built indentation. It also covers all
collection literals, comprehensions and glued string literals in
parentheses most of which were overlooked in the previous implementation.

python/src/com/jetbrains/python/PyElementTypes.java

@@ -121,8 +121,6 @@ public interface PyElementTypes {
   PyElementType SET_COMP_EXPRESSION = new PyElementType("SET_COMP_EXPRESSION", PySetCompExpressionImpl.class);
   PyElementType DICT_COMP_EXPRESSION = new PyElementType("DICT_COMP_EXPRESSION", PyDictCompExpressionImpl.class);
 
-  TokenSet LIST_LIKE_EXPRESSIONS = TokenSet.create(LIST_LITERAL_EXPRESSION, LIST_COMP_EXPRESSION, TUPLE_EXPRESSION);
-
   TokenSet STATEMENT_LISTS = TokenSet.create(STATEMENT_LIST);
 
   TokenSet BINARY_OPS = TokenSet.create(PyTokenTypes.OR_KEYWORD, PyTokenTypes.AND_KEYWORD, PyTokenTypes.LT, PyTokenTypes.GT,

python/src/com/jetbrains/python/formatter/PyBlock.java

@@ -17,7 +17,6 @@ package com.jetbrains.python.formatter;
 
 import com.intellij.formatting.*;
 import com.intellij.lang.ASTNode;
-import com.intellij.openapi.editor.Document;
 import com.intellij.openapi.util.Key;
 import com.intellij.openapi.util.TextRange;
 import com.intellij.psi.*;
@@ -892,6 +891,8 @@ public class PyBlock implements ASTBlock {
   @NotNull
   public ChildAttributes getChildAttributes(int newChildIndex) {
     int statementListsBelow = 0;
+    PyBlock incompleteAlignedBlock = null;
+
     if (newChildIndex > 0) {
       // always pass decision to a sane block from top level from file or definition
       if (myNode.getPsi() instanceof PyFile || myNode.getElementType() == PyTokenTypes.COLON) {
@@ -935,6 +936,19 @@ public class PyBlock implements ASTBlock {
         }
         lastChild = getLastNonSpaceChild(lastChild, true);
       }
+
+      if ((prevElt instanceof PySequenceExpression || prevElt instanceof PyComprehensionElement) &&
+          prevElt.getLastChild() instanceof PsiErrorElement) {
+        incompleteAlignedBlock = insertAfterBlock;
+      }
+      else if (prevElt instanceof PyParenthesizedExpression && prevElt.getLastChild() instanceof PsiErrorElement) {
+        final PyExpression contained = ((PyParenthesizedExpression)prevElt).getContainedExpression();
+        // In case of parenthesized string literal we can use both the alignment of the containing parenthesized expression
+        // and the literal itself, since it's not clear whether user is going to insert another string node or the closing parenthesis
+        if (contained instanceof PyTupleExpression || contained instanceof PyStringLiteralExpression) {
+          incompleteAlignedBlock = insertAfterBlock.getSubBlockByNode(contained.getNode());
+        }
+      }
     }
 
     // HACKETY-HACK
@@ -959,7 +973,7 @@ public class PyBlock implements ASTBlock {
 
 
     final Indent childIndent = getChildIndent(newChildIndex);
-    final Alignment childAlignment = getChildAlignment();
+    final Alignment childAlignment = incompleteAlignedBlock != null ? incompleteAlignedBlock.getChildAlignment() : getChildAlignment();
     return new ChildAttributes(childIndent, childAlignment);
   }
 
@@ -985,15 +999,30 @@ public class PyBlock implements ASTBlock {
 
   @Nullable
   private Alignment getChildAlignment() {
-    if (ourListElementTypes.contains(myNode.getElementType()) || myNode.getElementType() == PyElementTypes.SLICE_ITEM) {
+    // TODO merge it with needListAlignment(ASTNode)
+    final IElementType nodeType = myNode.getElementType();
+    if (ourListElementTypes.contains(nodeType) ||
+        nodeType == PyElementTypes.SLICE_ITEM ||
+        nodeType == PyElementTypes.STRING_LITERAL_EXPRESSION) {
       if (isInControlStatement()) {
         return null;
       }
-      if (myNode.getPsi() instanceof PyParameterList && !myContext.getSettings().ALIGN_MULTILINE_PARAMETERS) {
+      final PsiElement elem = myNode.getPsi();
+      if (elem instanceof PyParameterList && !myContext.getSettings().ALIGN_MULTILINE_PARAMETERS) {
         return null;
       }
-      if (myNode.getPsi() instanceof PyDictLiteralExpression) {
-        final PyKeyValueExpression lastElement = ArrayUtil.getLastElement(((PyDictLiteralExpression)myNode.getPsi()).getElements());
+      if ((elem instanceof PySequenceExpression || elem instanceof PyComprehensionElement) &&
+          !myContext.getPySettings().ALIGN_COLLECTIONS_AND_COMPREHENSIONS) {
+        return null;
+      }
+      if (elem instanceof PyParenthesizedExpression) {
+        final PyExpression parenthesized = ((PyParenthesizedExpression)elem).getContainedExpression();
+        if (parenthesized instanceof PyTupleExpression && !myContext.getPySettings().ALIGN_COLLECTIONS_AND_COMPREHENSIONS) {
+          return null;
+        }
+      }
+      if (elem instanceof PyDictLiteralExpression) {
+        final PyKeyValueExpression lastElement = ArrayUtil.getLastElement(((PyDictLiteralExpression)elem).getElements());
         if (lastElement == null || lastElement.getValue() == null /* incomplete */) {
           return null;
         }
@@ -1025,44 +1054,6 @@ public class PyBlock implements ASTBlock {
         return Indent.getNormalIndent();
       }
     }
-    else if (lastChild != null && PyElementTypes.LIST_LIKE_EXPRESSIONS.contains(lastChild.getElementType())) {
-      // handle pressing enter at the end of a list literal when there's no closing paren or bracket
-      final ASTNode lastLastChild = lastChild.getLastChildNode();
-      if (lastLastChild != null && lastLastChild.getPsi() instanceof PsiErrorElement) {
-        // we're at a place like this: [foo, ... bar, <caret>
-        // we'd rather align to foo. this may be not a multiple of tabs.
-        final PsiElement expr = lastChild.getPsi();
-        PsiElement exprItem = expr.getFirstChild();
-        boolean found = false;
-        while (exprItem != null) { // find a worthy element to align to
-          if (exprItem instanceof PyElement) {
-            found = true; // align to foo in "[foo,"
-            break;
-          }
-          if (exprItem instanceof PsiComment) {
-            found = true; // align to foo in "[ # foo,"
-            break;
-          }
-          exprItem = exprItem.getNextSibling();
-        }
-        if (found) {
-          final PsiDocumentManager docMgr = PsiDocumentManager.getInstance(exprItem.getProject());
-          final Document doc = docMgr.getDocument(exprItem.getContainingFile());
-          if (doc != null) {
-            int lineNum = doc.getLineNumber(exprItem.getTextOffset());
-            final int itemCol = exprItem.getTextOffset() - doc.getLineStartOffset(lineNum);
-            final PsiElement hereElt = getNode().getPsi();
-            lineNum = doc.getLineNumber(hereElt.getTextOffset());
-            final int nodeCol = hereElt.getTextOffset() - doc.getLineStartOffset(lineNum);
-            final int padding = itemCol - nodeCol;
-            if (padding > 0) { // negative is a syntax error,  but possible
-              return Indent.getSpaceIndent(padding);
-            }
-          }
-        }
-        return Indent.getContinuationIndent(); // a fallback
-      }
-    }
 
     if (afterNode != null && afterNode.getElementType() == PyElementTypes.KEY_VALUE_EXPRESSION) {
       final PyKeyValueExpression keyValue = (PyKeyValueExpression)afterNode.getPsi();

python/testData/editing/enterInIncompleteDictComprehension.after.py

@@ -0,0 +1,2 @@
+mapping = {foo: bar for
+           foo

python/testData/editing/enterInIncompleteDictComprehension.py

@@ -0,0 +1 @@
+mapping = {foo: bar for<caret>

python/testData/editing/enterInIncompleteDictLiteral.after.py

@@ -0,0 +1,2 @@
+d = {'foo': 'bar',
+     'baz'

python/testData/editing/enterInIncompleteDictLiteral.py

@@ -0,0 +1 @@
+d = {'foo': 'bar',<caret>

python/testData/editing/enterInIncompleteGluedStringLiteralInParentheses.after.py

@@ -0,0 +1,2 @@
+xs = ('foo'
+      'bar'

python/testData/editing/enterInIncompleteGluedStringLiteralInParentheses.py

@@ -0,0 +1 @@
+xs = ('foo'<caret>

python/testData/editing/enterInIncompleteListComprehension.after.py

@@ -0,0 +1,2 @@
+xs = [foo for
+      foo

python/testData/editing/enterInIncompleteListComprehension.py

@@ -0,0 +1 @@
+xs = [foo for<caret>

python/testData/editing/enterInIncompleteListLiteral.after.py

@@ -0,0 +1,2 @@
+xs = ['foo', 'bar',
+      'baz'

python/testData/editing/enterInIncompleteListLiteral.py

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

python/testData/editing/enterInIncompleteParenthesizedGenerator.after.py

@@ -0,0 +1,2 @@
+xs = (foo for
+      foo

python/testData/editing/enterInIncompleteParenthesizedGenerator.py

@@ -0,0 +1 @@
+xs = (foo for<caret>

python/testData/editing/enterInIncompleteSetComprehension.after.py

@@ -0,0 +1,2 @@
+xs = {foo for
+      foo

python/testData/editing/enterInIncompleteSetComprehension.py

@@ -0,0 +1 @@
+xs = {foo for<caret>

python/testData/editing/enterInIncompleteSetLiteral.after.py

@@ -0,0 +1,2 @@
+xs = {'foo', 'bar',
+      'baz'

python/testData/editing/enterInIncompleteSetLiteral.py

@@ -0,0 +1 @@
+xs = {'foo', 'bar',<caret>

python/testData/editing/enterInIncompleteTupleLiteral.after.py

@@ -0,0 +1,2 @@
+xs = ('foo', 'bar',
+      'baz'<caret>

python/testData/editing/enterInIncompleteTupleLiteral.py

@@ -0,0 +1 @@
+xs = ('foo', 'bar',<caret>

python/testSrc/com/jetbrains/python/PyEditingTest.java

@@ -365,7 +365,8 @@ public class PyEditingTest extends PyTestCase {
   public void testEndOfStringInParenth() {
     doTestEnter("print (\"foo\"<caret>\n" +
                 "    \"bar\")",
-                "print (\"foo\"\n\n" +
+                "print (\"foo\"\n" +
+                "       \n" +
                 "    \"bar\")");
   }
 
@@ -596,4 +597,49 @@ public class PyEditingTest extends PyTestCase {
   public void testIncompleteFunctionTypeComment() {
     doTypingTest('.');
   }
+
+  // PY-10972
+  public void testEnterInIncompleteTupleLiteral() {
+    doTypingTest("\n'baz'");
+  }
+
+  // PY-10972
+  public void testEnterInIncompleteListLiteral() {
+    doTypingTest("\n'baz'");
+  }
+
+  // PY-10972
+  public void testEnterInIncompleteSetLiteral() {
+    doTypingTest("\n'baz'");
+  }
+
+  // PY-10972
+  public void testEnterInIncompleteDictLiteral() {
+    doTypingTest("\n'baz'");
+  }
+
+  // PY-10972
+  public void testEnterInIncompleteGluedStringLiteralInParentheses() {
+    doTypingTest("\n'bar'");
+  }
+
+  // PY-10972
+  public void testEnterInIncompleteListComprehension() {
+    doTypingTest("\nfoo");
+  }
+
+  // PY-10972
+  public void testEnterInIncompleteSetComprehension() {
+    doTypingTest("\nfoo");
+  }
+
+  // PY-10972
+  public void testEnterInIncompleteDictComprehension() {
+    doTypingTest("\nfoo");
+  }
+
+  // PY-10972
+  public void testEnterInIncompleteParenthesizedGenerator() {
+    doTypingTest("\nfoo");
+  }
 }