option to turn off alignment in collections and comprehensions (PY-8516)

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

@@ -149,7 +149,7 @@ public class PyBlock implements ASTBlock {
       }
     }
 
-    if (parentType == PyElementTypes.LIST_LITERAL_EXPRESSION) {
+    if (parentType == PyElementTypes.LIST_LITERAL_EXPRESSION || parentType == PyElementTypes.LIST_COMP_EXPRESSION) {
       if (childType == PyTokenTypes.RBRACKET || childType == PyTokenTypes.LBRACKET) {
         childIndent = Indent.getNoneIndent();
       }
@@ -167,7 +167,8 @@ public class PyBlock implements ASTBlock {
                       : Indent.getNormalIndent();
       }
     }
-    else if (parentType == PyElementTypes.DICT_LITERAL_EXPRESSION || parentType == PyElementTypes.SET_LITERAL_EXPRESSION) {
+    else if (parentType == PyElementTypes.DICT_LITERAL_EXPRESSION || parentType == PyElementTypes.SET_LITERAL_EXPRESSION ||
+      parentType == PyElementTypes.SET_COMP_EXPRESSION || parentType == PyElementTypes.DICT_COMP_EXPRESSION) {
       if (childType == PyTokenTypes.RBRACE || !hasLineBreaksBefore(child, 1)) {
         childIndent = Indent.getNoneIndent();
       }
@@ -192,7 +193,7 @@ public class PyBlock implements ASTBlock {
         childIndent = Indent.getNormalIndent();
       }
     }
-    else if ((parentType == PyElementTypes.PARENTHESIZED_EXPRESSION || parentType == PyElementTypes.GENERATOR_EXPRESSION)) {
+    else if (parentType == PyElementTypes.PARENTHESIZED_EXPRESSION || parentType == PyElementTypes.GENERATOR_EXPRESSION) {
       if (childType == PyTokenTypes.RPAR || !hasLineBreaksBefore(child, 1)) {
         childIndent = Indent.getNoneIndent();
       }
@@ -295,7 +296,7 @@ public class PyBlock implements ASTBlock {
     if (child.getElementType() == PyTokenTypes.COMMA) {
       return false;
     }
-    return true;
+    return myContext.getPySettings().ALIGN_COLLECTIONS_AND_COMPREHENSIONS;
   }
 
   @Nullable

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

@@ -2,18 +2,22 @@ package com.jetbrains.python.formatter;
 
 import com.intellij.formatting.FormattingMode;
 import com.intellij.formatting.SpacingBuilder;
+import com.intellij.psi.codeStyle.CodeStyleSettings;
 import com.intellij.psi.codeStyle.CommonCodeStyleSettings;
+import com.jetbrains.python.PythonLanguage;
 
 /**
  * @author yole
  */
 public class PyBlockContext {
   private final CommonCodeStyleSettings mySettings;
+  private final PyCodeStyleSettings myPySettings;
   private final SpacingBuilder mySpacingBuilder;
   private final FormattingMode myMode;
 
-  public PyBlockContext(CommonCodeStyleSettings settings, SpacingBuilder builder, FormattingMode mode) {
-    mySettings = settings;
+  public PyBlockContext(CodeStyleSettings settings, SpacingBuilder builder, FormattingMode mode) {
+    mySettings = settings.getCommonSettings(PythonLanguage.getInstance());
+    myPySettings = settings.getCustomSettings(PyCodeStyleSettings.class);
     mySpacingBuilder = builder;
     myMode = mode;
   }
@@ -22,6 +26,10 @@ public class PyBlockContext {
     return mySettings;
   }
 
+  public PyCodeStyleSettings getPySettings() {
+    return myPySettings;
+  }
+
   public SpacingBuilder getSpacingBuilder() {
     return mySpacingBuilder;
   }

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

@@ -16,6 +16,8 @@ public class PyCodeStyleSettings extends CustomCodeStyleSettings {
 
   public int BLANK_LINES_AROUND_TOP_LEVEL_CLASSES_FUNCTIONS = 2;
 
+  public boolean ALIGN_COLLECTIONS_AND_COMPREHENSIONS = true;
+
   public PyCodeStyleSettings(CodeStyleSettings container) {
     super("Python", container);
   }

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

@@ -73,6 +73,8 @@ public class PyLanguageCodeStyleSettingsProvider extends LanguageCodeStyleSettin
                                    "WRAP_LONG_LINES",
                                    "ALIGN_MULTILINE_PARAMETERS",
                                    "ALIGN_MULTILINE_PARAMETERS_IN_CALLS");
+      consumer.showCustomOption(PyCodeStyleSettings.class, "ALIGN_COLLECTIONS_AND_COMPREHENSIONS", "Align when multiline",
+                                "Collections and Comprehensions");
     }
   }
 
@@ -115,7 +117,9 @@ public class PyLanguageCodeStyleSettingsProvider extends LanguageCodeStyleSettin
                                                 "long_parameter_2):\n" +
                                                 "    pass\n\n" +
                                                 "xyzzy('long_string_constant1',\n" +
-                                                "    'long_string_constant2')";
+                                                "    'long_string_constant2')\n" +
+                                                "attrs = [e.attr for e in\n" +
+                                                "    items]";
   @SuppressWarnings("FieldCanBeLocal")
   private static String INDENT_SETTINGS_PREVIEW = "def foo():\n" +
                                                   "    print 'bar'\n\n" +

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

@@ -35,8 +35,7 @@ public class PythonFormattingModelBuilder implements FormattingModelBuilderEx, C
       System.out.println("AST tree for " + element.getContainingFile().getName() + ":");
       printAST(fileNode, 0);
     }
-    final CommonCodeStyleSettings codeStyleSettings = settings.getCommonSettings(PythonLanguage.getInstance());
-    final PyBlockContext context = new PyBlockContext(codeStyleSettings, createSpacingBuilder(settings), mode);
+    final PyBlockContext context = new PyBlockContext(settings, createSpacingBuilder(settings), mode);
     final PyBlock block = new PyBlock(null, element.getNode(), null, Indent.getNoneIndent(), null, context);
     if (DUMP_FORMATTING_AST) {
       FormattingModelDumper.dumpFormattingModel(block, 2, System.out);

python/testData/formatter/indentInComprehensions.py

@@ -0,0 +1,2 @@
+attrs = [e.attr for e in
+items]

python/testData/formatter/indentInComprehensions_after.py

@@ -0,0 +1,2 @@
+attrs = [e.attr for e in
+    items]

python/testSrc/com/jetbrains/python/PyFormatterTest.java

@@ -255,6 +255,11 @@ public class PyFormatterTest extends PyTestCase {
     doTest();
   }
 
+  public void testIndentInComprehensions() {  // PY-8516
+    settings().getCustomSettings(PyCodeStyleSettings.class).ALIGN_COLLECTIONS_AND_COMPREHENSIONS = false;
+    doTest();
+  }
+
   public void testContinuationIndentForCallInStatementPart() {  // PY-8577
     doTest();
   }